diff --git a/README.md b/README.md index 627418a9..e4a8bfd4 100644 --- a/README.md +++ b/README.md @@ -28,9 +28,14 @@ DCP reduces context size through a compress tool and automatic cleanup. Your ses ### Compress -Compress is a tool exposed to your model that selects a conversation range and replaces it with a technical summary. You can think of this as a much smarter version of Opencode's compaction process. Instead of triggering statically when your session reaches its maximum context and on the entire coding session, Compress allows the model to pick when to activate based on task completion, and to only compress a subset of messages containing the completed task. This allows the summaries replacing the session content to be much more focused and precise than Opencode's native compaction. +Compress is a tool exposed to your model that replaces closed, stale conversation content with high-fidelity technical summaries. You can think of this as a much smarter version of Opencode's compaction process. Instead of triggering statically when your session reaches its maximum context and on the entire coding session, Compress allows the model to pick when to activate based on task completion, and to only compress the specific messages that are no longer needed verbatim. -When a new compression overlaps an earlier one, the earlier summary is nested inside the new one — so information is preserved through layers of compression rather than diluted away. Additionally, protected tool outputs (such as subagents and skills) and protected file patterns are always kept in compression summaries, ensuring that the most important information is never lost. You can also enable `protectUserMessages` to preserve your messages verbatim during compression, though note that large prompts (e.g. copy-pasting log files in the prompt) will then never be compressed away. +DCP supports two compression modes: + +- `range` mode compresses a contiguous span of conversation into one or more reusable block summaries. +- `message` mode is experimental and compresses individual raw messages independently, letting the model manage context much more surgically around closed work. + +In `range` mode, when a new compression overlaps an earlier one, the earlier summary is nested inside the new one so information is preserved through layers of compression rather than diluted away. In both modes, protected tool outputs (such as subagents and skills) and protected file patterns are kept in compression summaries, ensuring that the most important information is never lost. You can also enable `protectUserMessages` to preserve your messages verbatim during compression, though note that large prompts (e.g. copy-pasting log files in the prompt) will then never be compressed away. ### Deduplication @@ -99,6 +104,9 @@ Each level overrides the previous, so project settings take priority over global "protectedFilePatterns": [], // Unified context compression tool and behavior settings "compress": { + // Compression mode: "range" (compress spans into block summaries) + // or experimental "message" (compress individual raw messages) + "mode": "range", // Permission mode: "allow" (no prompt), "ask" (prompt), "deny" (tool not registered) "permission": "allow", // Show compression content in a chat notification @@ -133,8 +141,6 @@ Each level overrides the previous, so project settings take priority over global // Controls how likely compression is after user messages // ("strong" = more likely, "soft" = less likely) "nudgeForce": "soft", - // Flat tool schema: improves tool call reliability but uglier in the TUI - "flatSchema": false, // Tool names whose completed outputs are appended to the compression "protectedTools": [], // Preserve your messages during compression. @@ -149,10 +155,6 @@ Each level overrides the previous, so project settings take priority over global // Additional tools to protect from pruning "protectedTools": [], }, - // Prune write tool inputs when the file has been subsequently read - "supersedeWrites": { - "enabled": true, - }, // Prune tool inputs for errored tools after X turns "purgeErrors": { "enabled": true, @@ -176,16 +178,17 @@ DCP provides a `/dcp` slash command: - `/dcp stats` — Shows cumulative pruning statistics across all sessions. - `/dcp sweep` — Prunes all tools since the last user message. Accepts an optional count: `/dcp sweep 10` prunes the last 10 tools. Respects `commands.protectedTools`. - `/dcp manual [on|off]` — Toggle manual mode or set explicit state. When on, the AI will not autonomously use context management tools. -- `/dcp compress [focus]` — Trigger a single compress tool execution. Optional focus text directs what range to compress. +- `/dcp compress [focus]` — Trigger a single compress tool execution. Optional focus text directs what content to compress, following the active `compress.mode`. - `/dcp decompress ` — Restore a specific active compression by ID (for example `/dcp decompress 2`). Running without an argument shows available compression IDs, token sizes, and topics. - `/dcp recompress ` — Re-apply a user-decompressed compression by ID (for example `/dcp recompress 2`). Running without an argument shows recompressible IDs, token sizes, and topics. ### Prompt Overrides -DCP exposes five editable prompts: +DCP exposes six editable prompts: - `system` -- `compress` +- `compress-range` +- `compress-message` - `context-limit-nudge` - `turn-nudge` - `iteration-nudge` @@ -199,16 +202,16 @@ To customize behavior, add a file with the same name under an overrides director To reset an override, delete the matching file from your overrides directory. > [!NOTE] -> `compress` prompt changes apply after plugin restart because tool descriptions are registered at startup. +> `compress-range` and `compress-message` prompt changes apply after plugin restart because tool descriptions are registered at startup. ### Protected Tools By default, these tools are always protected from pruning: -`task`, `skill`, `todowrite`, `todoread`, `compress`, `batch`, `plan_enter`, `plan_exit` +`task`, `skill`, `todowrite`, `todoread`, `compress`, `batch`, `plan_enter`, `plan_exit`, `write`, `edit` The `protectedTools` arrays in `commands` and `strategies` add to this default list. -For the `compress` tool, `compress.protectedTools` ensures specific tool outputs are appended to the compressed summary. It defaults to an empty array `[]` but always inherently protects `task`, `skill`, `todowrite`, and `todoread`. +For the `compress` tool, `compress.protectedTools` ensures specific tool outputs are appended to the compressed summary. By default it includes `task`, `skill`, `todowrite`, and `todoread`. ## Impact on Prompt Caching diff --git a/dcp.schema.json b/dcp.schema.json index a927b870..23fe3642 100644 --- a/dcp.schema.json +++ b/dcp.schema.json @@ -69,7 +69,7 @@ "automaticStrategies": { "type": "boolean", "default": true, - "description": "When manual mode is enabled, keep automatic deduplication/supersede/purge strategies running" + "description": "When manual mode is enabled, keep automatic deduplication/purge strategies running" } }, "default": { @@ -128,6 +128,12 @@ "description": "Configuration for the unified compress tool", "additionalProperties": false, "properties": { + "mode": { + "type": "string", + "enum": ["range", "message"], + "default": "range", + "description": "Compression mode. 'range' compresses spans into block summaries, 'message' compresses individual raw messages." + }, "permission": { "type": "string", "enum": ["ask", "allow", "deny"], @@ -213,11 +219,6 @@ "default": "soft", "description": "Controls how likely compression is after user messages. 'strong' is more likely, 'soft' is less likely." }, - "flatSchema": { - "type": "boolean", - "default": false, - "description": "When true, the compress tool schema uses 4 flat string parameters (topic, startId, endId, summary) instead of the nested content object. This simplifies tool calls but changes TUI display." - }, "protectedTools": { "type": "array", "items": { @@ -231,6 +232,18 @@ "default": false, "description": "When enabled, your messages are never lost during compression" } + }, + "default": { + "mode": "range", + "permission": "allow", + "showCompression": false, + "maxContextLimit": 150000, + "minContextLimit": 50000, + "nudgeFrequency": 5, + "iterationNudgeThreshold": 15, + "nudgeForce": "soft", + "protectedTools": [], + "protectUserMessages": false } }, "strategies": { @@ -258,18 +271,6 @@ } } }, - "supersedeWrites": { - "type": "object", - "description": "Replace older write/edit outputs when new ones target the same file", - "additionalProperties": false, - "properties": { - "enabled": { - "type": "boolean", - "default": true, - "description": "Enable supersede writes strategy" - } - } - }, "purgeErrors": { "type": "object", "description": "Remove tool outputs that resulted in errors", diff --git a/index.ts b/index.ts index 86d541e7..eaff3c7e 100644 --- a/index.ts +++ b/index.ts @@ -1,5 +1,6 @@ import type { Plugin } from "@opencode-ai/plugin" import { getConfig } from "./lib/config" +import { createCompressMessageTool, createCompressRangeTool } from "./lib/compress" import { compressDisabledByOpencode, hasExplicitToolPermission, @@ -7,7 +8,6 @@ import { } from "./lib/host-permissions" import { Logger } from "./lib/logger" import { createSessionState } from "./lib/state" -import { createCompressTool } from "./lib/tools" import { PromptStore } from "./lib/prompts/store" import { createChatMessageTransformHandler, @@ -41,6 +41,14 @@ const plugin: Plugin = (async (ctx) => { strategies: config.strategies, }) + const compressToolContext = { + client: ctx.client, + state, + logger, + config, + prompts, + } + return { "experimental.chat.system.transform": createSystemPromptHandler( state, @@ -81,14 +89,10 @@ const plugin: Plugin = (async (ctx) => { ), tool: { ...(config.compress.permission !== "deny" && { - compress: createCompressTool({ - client: ctx.client, - state, - logger, - config, - workingDirectory: ctx.directory, - prompts, - }), + compress: + config.compress.mode === "message" + ? createCompressMessageTool(compressToolContext) + : createCompressRangeTool(compressToolContext), }), }, config: async (opencodeConfig) => { diff --git a/lib/commands/compression-targets.ts b/lib/commands/compression-targets.ts new file mode 100644 index 00000000..887ad53e --- /dev/null +++ b/lib/commands/compression-targets.ts @@ -0,0 +1,135 @@ +import type { CompressionBlock, PruneMessagesState } from "../state" + +export interface CompressionTarget { + displayId: number + runId: number + topic: string + compressedTokens: number + grouped: boolean + blocks: CompressionBlock[] +} + +function byBlockId(a: CompressionBlock, b: CompressionBlock): number { + return a.blockId - b.blockId +} + +function buildTarget(blocks: CompressionBlock[]): CompressionTarget { + const ordered = [...blocks].sort(byBlockId) + const first = ordered[0] + if (!first) { + throw new Error("Cannot build compression target from empty block list.") + } + + const grouped = first.mode === "message" + return { + displayId: first.blockId, + runId: first.runId, + topic: grouped ? first.batchTopic || first.topic : first.topic, + compressedTokens: ordered.reduce((total, block) => total + block.compressedTokens, 0), + grouped, + blocks: ordered, + } +} + +function groupMessageBlocks(blocks: CompressionBlock[]): CompressionTarget[] { + const grouped = new Map() + + for (const block of blocks) { + const existing = grouped.get(block.runId) + if (existing) { + existing.push(block) + continue + } + grouped.set(block.runId, [block]) + } + + return Array.from(grouped.values()).map(buildTarget) +} + +function splitTargets(blocks: CompressionBlock[]): CompressionTarget[] { + const messageBlocks: CompressionBlock[] = [] + const singleBlocks: CompressionBlock[] = [] + + for (const block of blocks) { + if (block.mode === "message") { + messageBlocks.push(block) + } else { + singleBlocks.push(block) + } + } + + const targets = [ + ...singleBlocks.map((block) => buildTarget([block])), + ...groupMessageBlocks(messageBlocks), + ] + return targets.sort((a, b) => a.displayId - b.displayId) +} + +export function getActiveCompressionTargets( + messagesState: PruneMessagesState, +): CompressionTarget[] { + const activeBlocks = Array.from(messagesState.activeBlockIds) + .map((blockId) => messagesState.blocksById.get(blockId)) + .filter((block): block is CompressionBlock => !!block && block.active) + + return splitTargets(activeBlocks) +} + +export function getRecompressibleCompressionTargets( + messagesState: PruneMessagesState, + availableMessageIds: Set, +): CompressionTarget[] { + const allBlocks = Array.from(messagesState.blocksById.values()).filter((block) => { + return availableMessageIds.has(block.compressMessageId) + }) + + const messageGroups = new Map() + const singleTargets: CompressionTarget[] = [] + + for (const block of allBlocks) { + if (block.mode === "message") { + const existing = messageGroups.get(block.runId) + if (existing) { + existing.push(block) + } else { + messageGroups.set(block.runId, [block]) + } + continue + } + + if (block.deactivatedByUser && !block.active) { + singleTargets.push(buildTarget([block])) + } + } + + for (const blocks of messageGroups.values()) { + if (blocks.some((block) => block.deactivatedByUser && !block.active)) { + singleTargets.push(buildTarget(blocks)) + } + } + + return singleTargets.sort((a, b) => a.displayId - b.displayId) +} + +export function resolveCompressionTarget( + messagesState: PruneMessagesState, + blockId: number, +): CompressionTarget | null { + const block = messagesState.blocksById.get(blockId) + if (!block) { + return null + } + + if (block.mode !== "message") { + return buildTarget([block]) + } + + const blocks = Array.from(messagesState.blocksById.values()).filter( + (candidate) => candidate.mode === "message" && candidate.runId === block.runId, + ) + if (blocks.length === 0) { + return null + } + + return buildTarget(blocks) +} diff --git a/lib/commands/decompress.ts b/lib/commands/decompress.ts index 0d67dca6..5957632c 100644 --- a/lib/commands/decompress.ts +++ b/lib/commands/decompress.ts @@ -6,6 +6,11 @@ import { getCurrentParams } from "../strategies/utils" import { saveSessionState } from "../state/persistence" import { sendIgnoredMessage } from "../ui/notification" import { formatTokenCount } from "../ui/utils" +import { + getActiveCompressionTargets, + resolveCompressionTarget, + type CompressionTarget, +} from "./compression-targets" export interface DecompressCommandContext { client: any @@ -31,13 +36,6 @@ function parseBlockIdArg(arg: string): number | null { return Number.isInteger(parsed) && parsed > 0 ? parsed : null } -function getAvailableBlocks(messagesState: PruneMessagesState): CompressionBlock[] { - return Array.from(messagesState.activeBlockIds) - .map((blockId) => messagesState.blocksById.get(blockId)) - .filter((block): block is CompressionBlock => !!block && block.active) - .sort((a, b) => a.blockId - b.blockId) -} - function findActiveParentBlockId( messagesState: PruneMessagesState, block: CompressionBlock, @@ -71,6 +69,20 @@ function findActiveParentBlockId( return null } +function findActiveAncestorBlockId( + messagesState: PruneMessagesState, + target: CompressionTarget, +): number | null { + for (const block of target.blocks) { + const activeAncestorBlockId = findActiveParentBlockId(messagesState, block) + if (activeAncestorBlockId !== null) { + return activeAncestorBlockId + } + } + + return null +} + function snapshotActiveMessages(messagesState: PruneMessagesState): Map { const activeMessages = new Map() for (const [messageId, entry] of messagesState.byMessageId) { @@ -82,14 +94,17 @@ function snapshotActiveMessages(messagesState: PruneMessagesState): Map 0) { const refs = reactivatedBlockIds.map((id) => String(id)).join(", ") lines.push(`Also restored nested compression(s): ${refs}.`) @@ -106,22 +121,25 @@ function formatDecompressMessage( return lines.join("\n") } -function formatAvailableBlocksMessage(availableBlocks: CompressionBlock[]): string { +function formatAvailableBlocksMessage(availableTargets: CompressionTarget[]): string { const lines: string[] = [] lines.push("Usage: /dcp decompress ") lines.push("") - if (availableBlocks.length === 0) { + if (availableTargets.length === 0) { lines.push("No compressions are available to restore.") return lines.join("\n") } lines.push("Available compressions:") - const entries = availableBlocks.map((block) => { - const topic = block.topic.replace(/\s+/g, " ").trim() || "(no topic)" - const label = `${block.blockId} (${formatTokenCount(block.compressedTokens)})` - return { label, topic } + const entries = availableTargets.map((target) => { + const topic = target.topic.replace(/\s+/g, " ").trim() || "(no topic)" + const label = `${target.displayId} (${formatTokenCount(target.compressedTokens)})` + const details = target.grouped + ? `Compression #${target.runId} - ${target.blocks.length} messages` + : `Compression #${target.runId}` + return { label, topic: `${details} - ${topic}` } }) const labelWidth = Math.max(...entries.map((entry) => entry.label.length)) + 4 @@ -153,8 +171,8 @@ export async function handleDecompressCommand(ctx: DecompressCommandContext): Pr const messagesState = state.prune.messages if (!targetArg) { - const availableBlocks = getAvailableBlocks(messagesState) - const message = formatAvailableBlocksMessage(availableBlocks) + const availableTargets = getActiveCompressionTargets(messagesState) + const message = formatAvailableBlocksMessage(availableTargets) await sendIgnoredMessage(client, sessionId, message, params, logger) return } @@ -171,8 +189,8 @@ export async function handleDecompressCommand(ctx: DecompressCommandContext): Pr return } - const targetBlock = messagesState.blocksById.get(targetBlockId) - if (!targetBlock) { + const target = resolveCompressionTarget(messagesState, targetBlockId) + if (!target) { await sendIgnoredMessage( client, sessionId, @@ -183,13 +201,14 @@ export async function handleDecompressCommand(ctx: DecompressCommandContext): Pr return } - if (!targetBlock.active) { - const activeAncestorBlockId = findActiveParentBlockId(messagesState, targetBlock) + const activeBlocks = target.blocks.filter((block) => block.active) + if (activeBlocks.length === 0) { + const activeAncestorBlockId = findActiveAncestorBlockId(messagesState, target) if (activeAncestorBlockId !== null) { await sendIgnoredMessage( client, sessionId, - `Compression ${targetBlockId} is inside compression ${activeAncestorBlockId}. Restore compression ${activeAncestorBlockId} first.`, + `Compression ${target.displayId} is inside compression ${activeAncestorBlockId}. Restore compression ${activeAncestorBlockId} first.`, params, logger, ) @@ -199,7 +218,7 @@ export async function handleDecompressCommand(ctx: DecompressCommandContext): Pr await sendIgnoredMessage( client, sessionId, - `Compression ${targetBlockId} is not active.`, + `Compression ${target.displayId} is not active.`, params, logger, ) @@ -208,11 +227,14 @@ export async function handleDecompressCommand(ctx: DecompressCommandContext): Pr const activeMessagesBefore = snapshotActiveMessages(messagesState) const activeBlockIdsBefore = new Set(messagesState.activeBlockIds) + const deactivatedAt = Date.now() - targetBlock.active = false - targetBlock.deactivatedByUser = true - targetBlock.deactivatedAt = Date.now() - targetBlock.deactivatedByBlockId = undefined + for (const block of target.blocks) { + block.active = false + block.deactivatedByUser = true + block.deactivatedAt = deactivatedAt + block.deactivatedByBlockId = undefined + } syncCompressionBlocks(state, logger, messages) @@ -236,7 +258,7 @@ export async function handleDecompressCommand(ctx: DecompressCommandContext): Pr await saveSessionState(state, logger) const message = formatDecompressMessage( - targetBlockId, + target, restoredMessageCount, restoredTokens, reactivatedBlockIds, @@ -244,7 +266,8 @@ export async function handleDecompressCommand(ctx: DecompressCommandContext): Pr await sendIgnoredMessage(client, sessionId, message, params, logger) logger.info("Decompress command completed", { - targetBlockId, + targetBlockId: target.displayId, + targetRunId: target.runId, restoredMessageCount, restoredTokens, reactivatedBlockIds, diff --git a/lib/commands/manual.ts b/lib/commands/manual.ts index df1a1428..aa8bb1dd 100644 --- a/lib/commands/manual.ts +++ b/lib/commands/manual.ts @@ -21,14 +21,20 @@ const MANUAL_MODE_OFF = "Manual mode is now OFF." const COMPRESS_TRIGGER_PROMPT = [ "", "Manual mode trigger received. You must now use the compress tool.", - "Find the most significant completed section of the conversation that can be compressed into a high-fidelity technical summary.", - "Choose safe boundaries and preserve all critical implementation details.", - "Return after compress with a brief explanation of what range was compressed.", + "Find the most significant completed conversation content that can be compressed into a high-fidelity technical summary.", + "Follow the active compress mode, preserve all critical implementation details, and choose safe targets.", + "Return after compress with a brief explanation of what content was compressed.", ].join("\n\n") -function getTriggerPrompt(tool: "compress", state: SessionState, userFocus?: string): string { +function getTriggerPrompt( + tool: "compress", + state: SessionState, + config: PluginConfig, + userFocus?: string, +): string { const base = COMPRESS_TRIGGER_PROMPT - const compressedBlockGuidance = buildCompressedBlockGuidance(state) + const compressedBlockGuidance = + config.compress.mode === "message" ? "" : buildCompressedBlockGuidance(state) const sections = [base, compressedBlockGuidance] if (userFocus && userFocus.trim().length > 0) { @@ -78,5 +84,5 @@ export async function handleManualTriggerCommand( tool: "compress", userFocus?: string, ): Promise { - return getTriggerPrompt(tool, ctx.state, userFocus) + return getTriggerPrompt(tool, ctx.state, ctx.config, userFocus) } diff --git a/lib/commands/recompress.ts b/lib/commands/recompress.ts index 4b39978b..eda5879b 100644 --- a/lib/commands/recompress.ts +++ b/lib/commands/recompress.ts @@ -1,11 +1,16 @@ import type { Logger } from "../logger" -import type { CompressionBlock, PruneMessagesState, SessionState, WithParts } from "../state" +import type { PruneMessagesState, SessionState, WithParts } from "../state" import { syncCompressionBlocks } from "../messages" import { parseBlockRef } from "../message-ids" import { getCurrentParams } from "../strategies/utils" import { saveSessionState } from "../state/persistence" import { sendIgnoredMessage } from "../ui/notification" import { formatTokenCount } from "../ui/utils" +import { + getRecompressibleCompressionTargets, + resolveCompressionTarget, + type CompressionTarget, +} from "./compression-targets" export interface RecompressCommandContext { client: any @@ -31,20 +36,6 @@ function parseBlockIdArg(arg: string): number | null { return Number.isInteger(parsed) && parsed > 0 ? parsed : null } -function getRecompressibleBlocks( - messagesState: PruneMessagesState, - availableMessageIds: Set, -): CompressionBlock[] { - return Array.from(messagesState.blocksById.values()) - .filter( - (block) => - block.deactivatedByUser && - !block.active && - availableMessageIds.has(block.compressMessageId), - ) - .sort((a, b) => a.blockId - b.blockId) -} - function snapshotActiveMessages(messagesState: PruneMessagesState): Set { const activeMessages = new Set() for (const [messageId, entry] of messagesState.byMessageId) { @@ -56,14 +47,17 @@ function snapshotActiveMessages(messagesState: PruneMessagesState): Set } function formatRecompressMessage( - targetBlockId: number, + target: CompressionTarget, recompressedMessageCount: number, recompressedTokens: number, deactivatedBlockIds: number[], ): string { const lines: string[] = [] - lines.push(`Re-applied compression ${targetBlockId}.`) + lines.push(`Re-applied compression ${target.displayId}.`) + if (target.runId !== target.displayId || target.grouped) { + lines.push(`Tool call label: Compression #${target.runId}.`) + } if (deactivatedBlockIds.length > 0) { const refs = deactivatedBlockIds.map((id) => String(id)).join(", ") lines.push(`Also re-compressed nested compression(s): ${refs}.`) @@ -80,22 +74,25 @@ function formatRecompressMessage( return lines.join("\n") } -function formatAvailableBlocksMessage(availableBlocks: CompressionBlock[]): string { +function formatAvailableBlocksMessage(availableTargets: CompressionTarget[]): string { const lines: string[] = [] lines.push("Usage: /dcp recompress ") lines.push("") - if (availableBlocks.length === 0) { + if (availableTargets.length === 0) { lines.push("No user-decompressed blocks are available to re-compress.") return lines.join("\n") } - lines.push("Available user-decompressed blocks:") - const entries = availableBlocks.map((block) => { - const topic = block.topic.replace(/\s+/g, " ").trim() || "(no topic)" - const label = `${block.blockId} (${formatTokenCount(block.compressedTokens)})` - return { label, topic } + lines.push("Available user-decompressed compressions:") + const entries = availableTargets.map((target) => { + const topic = target.topic.replace(/\s+/g, " ").trim() || "(no topic)" + const label = `${target.displayId} (${formatTokenCount(target.compressedTokens)})` + const details = target.grouped + ? `Compression #${target.runId} - ${target.blocks.length} messages` + : `Compression #${target.runId}` + return { label, topic: `${details} - ${topic}` } }) const labelWidth = Math.max(...entries.map((entry) => entry.label.length)) + 4 @@ -128,8 +125,11 @@ export async function handleRecompressCommand(ctx: RecompressCommandContext): Pr const availableMessageIds = new Set(messages.map((msg) => msg.info.id)) if (!targetArg) { - const availableBlocks = getRecompressibleBlocks(messagesState, availableMessageIds) - const message = formatAvailableBlocksMessage(availableBlocks) + const availableTargets = getRecompressibleCompressionTargets( + messagesState, + availableMessageIds, + ) + const message = formatAvailableBlocksMessage(availableTargets) await sendIgnoredMessage(client, sessionId, message, params, logger) return } @@ -146,8 +146,8 @@ export async function handleRecompressCommand(ctx: RecompressCommandContext): Pr return } - const targetBlock = messagesState.blocksById.get(targetBlockId) - if (!targetBlock) { + const target = resolveCompressionTarget(messagesState, targetBlockId) + if (!target) { await sendIgnoredMessage( client, sessionId, @@ -158,21 +158,21 @@ export async function handleRecompressCommand(ctx: RecompressCommandContext): Pr return } - if (!availableMessageIds.has(targetBlock.compressMessageId)) { + if (target.blocks.some((block) => !availableMessageIds.has(block.compressMessageId))) { await sendIgnoredMessage( client, sessionId, - `Compression ${targetBlockId} can no longer be re-applied because its origin message is no longer in this session.`, + `Compression ${target.displayId} can no longer be re-applied because its origin message is no longer in this session.`, params, logger, ) return } - if (!targetBlock.deactivatedByUser) { - const message = targetBlock.active - ? `Compression ${targetBlockId} is already active.` - : `Compression ${targetBlockId} is not user-decompressed.` + if (!target.blocks.some((block) => block.deactivatedByUser)) { + const message = target.blocks.some((block) => block.active) + ? `Compression ${target.displayId} is already active.` + : `Compression ${target.displayId} is not user-decompressed.` await sendIgnoredMessage(client, sessionId, message, params, logger) return } @@ -180,9 +180,11 @@ export async function handleRecompressCommand(ctx: RecompressCommandContext): Pr const activeMessagesBefore = snapshotActiveMessages(messagesState) const activeBlockIdsBefore = new Set(messagesState.activeBlockIds) - targetBlock.deactivatedByUser = false - targetBlock.deactivatedAt = undefined - targetBlock.deactivatedByBlockId = undefined + for (const block of target.blocks) { + block.deactivatedByUser = false + block.deactivatedAt = undefined + block.deactivatedByBlockId = undefined + } syncCompressionBlocks(state, logger, messages) @@ -205,7 +207,7 @@ export async function handleRecompressCommand(ctx: RecompressCommandContext): Pr await saveSessionState(state, logger) const message = formatRecompressMessage( - targetBlockId, + target, recompressedMessageCount, recompressedTokens, deactivatedBlockIds, @@ -213,7 +215,8 @@ export async function handleRecompressCommand(ctx: RecompressCommandContext): Pr await sendIgnoredMessage(client, sessionId, message, params, logger) logger.info("Recompress command completed", { - targetBlockId, + targetBlockId: target.displayId, + targetRunId: target.runId, recompressedMessageCount, recompressedTokens, deactivatedBlockIds, diff --git a/lib/compress/index.ts b/lib/compress/index.ts new file mode 100644 index 00000000..bdb7f2eb --- /dev/null +++ b/lib/compress/index.ts @@ -0,0 +1,3 @@ +export { ToolContext } from "./types" +export { createCompressMessageTool } from "./message" +export { createCompressRangeTool } from "./range" diff --git a/lib/compress/message-utils.ts b/lib/compress/message-utils.ts new file mode 100644 index 00000000..90f131bb --- /dev/null +++ b/lib/compress/message-utils.ts @@ -0,0 +1,190 @@ +import type { PluginConfig } from "../config" +import type { SessionState } from "../state" +import { parseBoundaryId } from "../message-ids" +import { isIgnoredUserMessage, isProtectedUserMessage } from "../messages/utils" +import { resolveAnchorMessageId, resolveBoundaryIds, resolveSelection } from "./search" +import { COMPRESSED_BLOCK_HEADER } from "./state" +import type { + CompressMessageEntry, + CompressMessageToolArgs, + ResolvedMessageCompression, + ResolvedMessageCompressionsResult, + SearchContext, +} from "./types" + +class SoftIssue extends Error {} + +export function validateArgs(args: CompressMessageToolArgs): void { + if (typeof args.topic !== "string" || args.topic.trim().length === 0) { + throw new Error("topic is required and must be a non-empty string") + } + + if (!Array.isArray(args.content) || args.content.length === 0) { + throw new Error("content is required and must be a non-empty array") + } + + for (let index = 0; index < args.content.length; index++) { + const entry = args.content[index] + const prefix = `content[${index}]` + + if (typeof entry?.messageId !== "string" || entry.messageId.trim().length === 0) { + throw new Error(`${prefix}.messageId is required and must be a non-empty string`) + } + + if (typeof entry?.topic !== "string" || entry.topic.trim().length === 0) { + throw new Error(`${prefix}.topic is required and must be a non-empty string`) + } + + if (typeof entry?.summary !== "string" || entry.summary.trim().length === 0) { + throw new Error(`${prefix}.summary is required and must be a non-empty string`) + } + } +} + +export function formatResult(processedCount: number, skippedIssues: string[]): string { + const messageNoun = processedCount === 1 ? "message" : "messages" + const processedText = + processedCount > 0 + ? `Compressed ${processedCount} ${messageNoun} into ${COMPRESSED_BLOCK_HEADER}.` + : "Compressed 0 messages." + + if (skippedIssues.length === 0) { + return processedText + } + + const issueNoun = skippedIssues.length === 1 ? "issue" : "issues" + const issueLines = skippedIssues.map((issue) => `- ${issue}`).join("\n") + return `${processedText}\nSkipped ${skippedIssues.length} ${issueNoun}:\n${issueLines}` +} + +export function formatIssues(skippedIssues: string[]): string { + const issueNoun = skippedIssues.length === 1 ? "issue" : "issues" + const issueLines = skippedIssues.map((issue) => `- ${issue}`).join("\n") + return `Unable to compress any messages. Found ${skippedIssues.length} ${issueNoun}:\n${issueLines}` +} + +export function resolveMessages( + args: CompressMessageToolArgs, + searchContext: SearchContext, + state: SessionState, + config: PluginConfig, +): ResolvedMessageCompressionsResult { + const issues: string[] = [] + const plans: ResolvedMessageCompression[] = [] + const seenMessageIds = new Set() + + for (const entry of args.content) { + const normalizedMessageId = entry.messageId.trim() + if (seenMessageIds.has(normalizedMessageId)) { + issues.push( + `messageId ${normalizedMessageId} was selected more than once in this batch.`, + ) + continue + } + + try { + const plan = resolveMessage( + { + ...entry, + messageId: normalizedMessageId, + }, + searchContext, + state, + config, + ) + seenMessageIds.add(plan.entry.messageId) + plans.push(plan) + } catch (error: any) { + if (error instanceof SoftIssue) { + issues.push(error.message) + continue + } + + throw error + } + } + + return { + plans, + skippedIssues: issues, + } +} + +function resolveMessage( + entry: CompressMessageEntry, + searchContext: SearchContext, + state: SessionState, + config: PluginConfig, +): ResolvedMessageCompression { + if (entry.messageId.toUpperCase() === "BLOCKED") { + throw new SoftIssue( + "messageId BLOCKED refers to a protected message and cannot be compressed.", + ) + } + + const parsed = parseBoundaryId(entry.messageId) + + if (!parsed) { + throw new Error( + `messageId ${entry.messageId} is invalid. Use an injected raw message ID of the form mNNNN.`, + ) + } + + if (parsed.kind === "compressed-block") { + throw new SoftIssue( + `messageId ${entry.messageId} is invalid here. Block IDs like bN are not allowed; use an mNNNN message ID instead.`, + ) + } + + const messageId = state.messageIds.byRef.get(parsed.ref) + const rawMessage = messageId ? searchContext.rawMessagesById.get(messageId) : undefined + const hasBoundary = + !!rawMessage && + !!messageId && + searchContext.rawIndexById.has(messageId) && + !(rawMessage.info.role === "user" && isIgnoredUserMessage(rawMessage)) + if (!hasBoundary) { + throw new SoftIssue( + `messageId ${parsed.ref} is not available in the current conversation context. Choose an injected mNNNN ID visible in context.`, + ) + } + + const { startReference, endReference } = resolveBoundaryIds( + searchContext, + state, + parsed.ref, + parsed.ref, + ) + const selection = resolveSelection(searchContext, startReference, endReference) + const rawMessageId = selection.messageIds[0] + + if (!rawMessageId) { + throw new Error(`messageId ${parsed.ref} could not be resolved to a raw message.`) + } + + const message = searchContext.rawMessagesById.get(rawMessageId) + if (!message) { + throw new Error(`messageId ${parsed.ref} is not available in the current conversation.`) + } + + if (isProtectedUserMessage(config, message)) { + throw new SoftIssue( + `messageId ${parsed.ref} refers to a protected message and cannot be compressed.`, + ) + } + + const pruneEntry = state.prune.messages.byMessageId.get(rawMessageId) + if (pruneEntry && pruneEntry.activeBlockIds.length > 0) { + throw new Error(`messageId ${parsed.ref} is already part of an active compression.`) + } + + return { + entry: { + messageId: parsed.ref, + topic: entry.topic, + summary: entry.summary, + }, + selection, + anchorMessageId: resolveAnchorMessageId(startReference), + } +} diff --git a/lib/compress/message.ts b/lib/compress/message.ts new file mode 100644 index 00000000..a7b91a77 --- /dev/null +++ b/lib/compress/message.ts @@ -0,0 +1,131 @@ +import { tool } from "@opencode-ai/plugin" +import type { ToolContext } from "./types" +import { countTokens } from "../strategies/utils" +import { MESSAGE_FORMAT_OVERLAY } from "../prompts/internal-overlays" +import { formatIssues, formatResult, resolveMessages, validateArgs } from "./message-utils" +import { finalizeSession, prepareSession, type NotificationEntry } from "./pipeline" +import { appendProtectedTools } from "./protected-content" +import { + allocateBlockId, + allocateRunId, + applyCompressionState, + wrapCompressedSummary, +} from "./state" +import type { CompressMessageToolArgs } from "./types" + +function buildSchema() { + return { + topic: tool.schema + .string() + .describe( + "Short label (3-5 words) for the overall batch - e.g., 'Closed Research Notes'", + ), + content: tool.schema + .array( + tool.schema.object({ + messageId: tool.schema + .string() + .describe("Raw message ID to compress (e.g. m0001)"), + topic: tool.schema + .string() + .describe("Short label (3-5 words) for this one message summary"), + summary: tool.schema + .string() + .describe("Complete technical summary replacing that one message"), + }), + ) + .describe("Batch of individual message summaries to create in one tool call"), + } +} + +export function createCompressMessageTool(ctx: ToolContext): ReturnType { + ctx.prompts.reload() + const runtimePrompts = ctx.prompts.getRuntimePrompts() + + return tool({ + description: runtimePrompts.compressMessage + MESSAGE_FORMAT_OVERLAY, + args: buildSchema(), + async execute(args, toolCtx) { + const input = args as CompressMessageToolArgs + validateArgs(input) + + const { rawMessages, searchContext } = await prepareSession( + ctx, + toolCtx, + `Compress Message: ${input.topic}`, + ) + const { plans, skippedIssues } = resolveMessages( + input, + searchContext, + ctx.state, + ctx.config, + ) + + if (plans.length === 0 && skippedIssues.length > 0) { + throw new Error(formatIssues(skippedIssues)) + } + + const notifications: NotificationEntry[] = [] + + const preparedPlans: Array<{ + plan: (typeof plans)[number] + summaryWithTools: string + }> = [] + + for (const plan of plans) { + const summaryWithTools = await appendProtectedTools( + ctx.client, + ctx.state, + ctx.config.experimental.allowSubAgents, + plan.entry.summary, + plan.selection, + searchContext, + ctx.config.compress.protectedTools, + ctx.config.protectedFilePatterns, + ) + + preparedPlans.push({ + plan, + summaryWithTools, + }) + } + + const runId = allocateRunId(ctx.state) + + for (const { plan, summaryWithTools } of preparedPlans) { + const blockId = allocateBlockId(ctx.state) + const storedSummary = wrapCompressedSummary(blockId, summaryWithTools) + const summaryTokens = countTokens(storedSummary) + + applyCompressionState( + ctx.state, + { + topic: plan.entry.topic, + batchTopic: input.topic, + startId: plan.entry.messageId, + endId: plan.entry.messageId, + mode: "message", + runId, + compressMessageId: toolCtx.messageID, + }, + plan.selection, + plan.anchorMessageId, + blockId, + storedSummary, + [], + ) + + notifications.push({ + blockId, + runId, + summary: summaryWithTools, + summaryTokens, + }) + } + + await finalizeSession(ctx, toolCtx, rawMessages, notifications, input.topic) + + return formatResult(plans.length, skippedIssues) + }, + }) +} diff --git a/lib/compress/pipeline.ts b/lib/compress/pipeline.ts new file mode 100644 index 00000000..be676de0 --- /dev/null +++ b/lib/compress/pipeline.ts @@ -0,0 +1,106 @@ +import type { WithParts } from "../state" +import { ensureSessionInitialized } from "../state" +import { saveSessionState } from "../state/persistence" +import { assignMessageRefs } from "../message-ids" +import { isIgnoredUserMessage } from "../messages/utils" +import { deduplicate, purgeErrors } from "../strategies" +import { getCurrentParams, getCurrentTokenUsage } from "../strategies/utils" +import { sendCompressNotification } from "../ui/notification" +import type { ToolContext } from "./types" +import { buildSearchContext, fetchSessionMessages } from "./search" +import type { SearchContext } from "./types" + +interface RunContext { + ask(input: { + permission: string + patterns: string[] + always: string[] + metadata: Record + }): Promise + metadata(input: { title: string }): void + sessionID: string +} + +export interface NotificationEntry { + blockId: number + runId: number + summary: string + summaryTokens: number +} + +export interface PreparedSession { + rawMessages: WithParts[] + searchContext: SearchContext +} + +export async function prepareSession( + ctx: ToolContext, + toolCtx: RunContext, + title: string, +): Promise { + if (ctx.state.manualMode && ctx.state.manualMode !== "compress-pending") { + throw new Error( + "Manual mode: compress blocked. Do not retry until `` appears in user context.", + ) + } + + await toolCtx.ask({ + permission: "compress", + patterns: ["*"], + always: ["*"], + metadata: {}, + }) + + toolCtx.metadata({ title }) + + const rawMessages = await fetchSessionMessages(ctx.client, toolCtx.sessionID) + + await ensureSessionInitialized( + ctx.client, + ctx.state, + toolCtx.sessionID, + ctx.logger, + rawMessages, + ctx.config.manualMode.enabled, + ) + + assignMessageRefs(ctx.state, rawMessages) + + deduplicate(ctx.state, ctx.logger, ctx.config, rawMessages) + purgeErrors(ctx.state, ctx.logger, ctx.config, rawMessages) + + return { + rawMessages, + searchContext: buildSearchContext(ctx.state, rawMessages), + } +} + +export async function finalizeSession( + ctx: ToolContext, + toolCtx: RunContext, + rawMessages: WithParts[], + entries: NotificationEntry[], + batchTopic: string | undefined, +): Promise { + ctx.state.manualMode = ctx.state.manualMode ? "active" : false + await saveSessionState(ctx.state, ctx.logger) + + const params = getCurrentParams(ctx.state, rawMessages, ctx.logger) + const totalSessionTokens = getCurrentTokenUsage(rawMessages) + const sessionMessageIds = rawMessages + .filter((msg) => !(msg.info.role === "user" && isIgnoredUserMessage(msg))) + .map((msg) => msg.info.id) + + await sendCompressNotification( + ctx.client, + ctx.logger, + ctx.config, + ctx.state, + toolCtx.sessionID, + entries, + batchTopic, + totalSessionTokens, + sessionMessageIds, + params, + ) +} diff --git a/lib/compress/protected-content.ts b/lib/compress/protected-content.ts new file mode 100644 index 00000000..3db4cf0a --- /dev/null +++ b/lib/compress/protected-content.ts @@ -0,0 +1,154 @@ +import type { SessionState } from "../state" +import { isIgnoredUserMessage } from "../messages/utils" +import { + getFilePathsFromParameters, + isFilePathProtected, + isToolNameProtected, +} from "../protected-patterns" +import { + buildSubagentResultText, + getSubAgentId, + mergeSubagentResult, +} from "../subagents/subagent-results" +import { fetchSessionMessages } from "./search" +import type { SearchContext, SelectionResolution } from "./types" + +export function appendProtectedUserMessages( + summary: string, + selection: SelectionResolution, + searchContext: SearchContext, + state: SessionState, + enabled: boolean, +): string { + if (!enabled) return summary + + const userTexts: string[] = [] + + for (const messageId of selection.messageIds) { + const existingCompressionEntry = state.prune.messages.byMessageId.get(messageId) + if (existingCompressionEntry && existingCompressionEntry.activeBlockIds.length > 0) { + continue + } + + const message = searchContext.rawMessagesById.get(messageId) + if (!message) continue + if (message.info.role !== "user") continue + if (isIgnoredUserMessage(message)) continue + + const parts = Array.isArray(message.parts) ? message.parts : [] + for (const part of parts) { + if (part.type === "text" && typeof part.text === "string" && part.text.trim()) { + userTexts.push(part.text) + break + } + } + } + + if (userTexts.length === 0) { + return summary + } + + const heading = "\n\nThe following user messages were sent in this conversation verbatim:" + const body = userTexts.map((text) => `\n${text}`).join("") + return summary + heading + body +} + +export async function appendProtectedTools( + client: any, + state: SessionState, + allowSubAgents: boolean, + summary: string, + selection: SelectionResolution, + searchContext: SearchContext, + protectedTools: string[], + protectedFilePatterns: string[] = [], +): Promise { + const protectedOutputs: string[] = [] + + for (const messageId of selection.messageIds) { + const existingCompressionEntry = state.prune.messages.byMessageId.get(messageId) + if (existingCompressionEntry && existingCompressionEntry.activeBlockIds.length > 0) { + continue + } + + const message = searchContext.rawMessagesById.get(messageId) + if (!message) continue + + const parts = Array.isArray(message.parts) ? message.parts : [] + for (const part of parts) { + if (part.type === "tool" && part.callID) { + let isToolProtected = isToolNameProtected(part.tool, protectedTools) + + if (!isToolProtected && protectedFilePatterns.length > 0) { + const filePaths = getFilePathsFromParameters(part.tool, part.state?.input) + if (isFilePathProtected(filePaths, protectedFilePatterns)) { + isToolProtected = true + } + } + + if (isToolProtected) { + const title = `Tool: ${part.tool}` + let output = "" + + if (part.state?.status === "completed" && part.state?.output) { + output = + typeof part.state.output === "string" + ? part.state.output + : JSON.stringify(part.state.output) + } + + if ( + allowSubAgents && + part.tool === "task" && + part.state?.status === "completed" && + typeof part.state?.output === "string" + ) { + const cachedSubAgentResult = state.subAgentResultCache.get(part.callID) + + if (cachedSubAgentResult !== undefined) { + if (cachedSubAgentResult) { + output = mergeSubagentResult( + part.state.output, + cachedSubAgentResult, + ) + } + } else { + const subAgentSessionId = getSubAgentId(part) + if (subAgentSessionId) { + let subAgentResultText = "" + try { + const subAgentMessages = await fetchSessionMessages( + client, + subAgentSessionId, + ) + subAgentResultText = buildSubagentResultText(subAgentMessages) + } catch { + subAgentResultText = "" + } + + if (subAgentResultText) { + state.subAgentResultCache.set(part.callID, subAgentResultText) + output = mergeSubagentResult( + part.state.output, + subAgentResultText, + ) + } + } + } + } + + if (output) { + protectedOutputs.push(`\n### ${title}\n${output}`) + } + } + } + } + } + + if (protectedOutputs.length === 0) { + return summary + } + + const heading = "\n\nThe following protected tools were used in this conversation as well:" + return summary + heading + protectedOutputs.join("") +} diff --git a/lib/compress/range-utils.ts b/lib/compress/range-utils.ts new file mode 100644 index 00000000..7aa8dbc9 --- /dev/null +++ b/lib/compress/range-utils.ts @@ -0,0 +1,308 @@ +import type { CompressionBlock, SessionState } from "../state" +import { resolveAnchorMessageId, resolveBoundaryIds, resolveSelection } from "./search" +import type { + BoundaryReference, + CompressRangeToolArgs, + InjectedSummaryResult, + ParsedBlockPlaceholder, + ResolvedRangeCompression, + SearchContext, +} from "./types" + +const BLOCK_PLACEHOLDER_REGEX = /\(b(\d+)\)|\{block_(\d+)\}/gi + +export function validateArgs(args: CompressRangeToolArgs): void { + if (typeof args.topic !== "string" || args.topic.trim().length === 0) { + throw new Error("topic is required and must be a non-empty string") + } + + if (!Array.isArray(args.content) || args.content.length === 0) { + throw new Error("content is required and must be a non-empty array") + } + + for (let index = 0; index < args.content.length; index++) { + const entry = args.content[index] + const prefix = `content[${index}]` + + if (typeof entry?.startId !== "string" || entry.startId.trim().length === 0) { + throw new Error(`${prefix}.startId is required and must be a non-empty string`) + } + + if (typeof entry?.endId !== "string" || entry.endId.trim().length === 0) { + throw new Error(`${prefix}.endId is required and must be a non-empty string`) + } + + if (typeof entry?.summary !== "string" || entry.summary.trim().length === 0) { + throw new Error(`${prefix}.summary is required and must be a non-empty string`) + } + } +} + +export function resolveRanges( + args: CompressRangeToolArgs, + searchContext: SearchContext, + state: SessionState, +): ResolvedRangeCompression[] { + return args.content.map((entry, index) => { + const normalizedEntry = { + startId: entry.startId.trim(), + endId: entry.endId.trim(), + summary: entry.summary, + } + + const { startReference, endReference } = resolveBoundaryIds( + searchContext, + state, + normalizedEntry.startId, + normalizedEntry.endId, + ) + const selection = resolveSelection(searchContext, startReference, endReference) + + return { + index, + entry: normalizedEntry, + selection, + anchorMessageId: resolveAnchorMessageId(startReference), + } + }) +} + +export function validateNonOverlapping(plans: ResolvedRangeCompression[]): void { + const sortedPlans = [...plans].sort( + (left, right) => + left.selection.startReference.rawIndex - right.selection.startReference.rawIndex || + left.selection.endReference.rawIndex - right.selection.endReference.rawIndex || + left.index - right.index, + ) + + const issues: string[] = [] + + for (let index = 1; index < sortedPlans.length; index++) { + const previous = sortedPlans[index - 1] + const current = sortedPlans[index] + if (!previous || !current) { + continue + } + + if (current.selection.startReference.rawIndex > previous.selection.endReference.rawIndex) { + continue + } + + issues.push( + `content[${previous.index}] (${previous.entry.startId}..${previous.entry.endId}) overlaps content[${current.index}] (${current.entry.startId}..${current.entry.endId}). Overlapping ranges cannot be compressed in the same batch.`, + ) + } + + if (issues.length > 0) { + throw new Error( + issues.length === 1 ? issues[0] : issues.map((issue) => `- ${issue}`).join("\n"), + ) + } +} + +export function parseBlockPlaceholders(summary: string): ParsedBlockPlaceholder[] { + const placeholders: ParsedBlockPlaceholder[] = [] + const regex = new RegExp(BLOCK_PLACEHOLDER_REGEX) + + let match: RegExpExecArray | null + while ((match = regex.exec(summary)) !== null) { + const full = match[0] + const blockIdPart = match[1] || match[2] + const parsed = Number.parseInt(blockIdPart, 10) + if (!Number.isInteger(parsed)) { + continue + } + + placeholders.push({ + raw: full, + blockId: parsed, + startIndex: match.index, + endIndex: match.index + full.length, + }) + } + + return placeholders +} + +export function validateSummaryPlaceholders( + placeholders: ParsedBlockPlaceholder[], + requiredBlockIds: number[], + startReference: BoundaryReference, + endReference: BoundaryReference, + summaryByBlockId: Map, +): number[] { + const boundaryOptionalIds = new Set() + if (startReference.kind === "compressed-block") { + if (startReference.blockId === undefined) { + throw new Error("Failed to map boundary matches back to raw messages") + } + boundaryOptionalIds.add(startReference.blockId) + } + if (endReference.kind === "compressed-block") { + if (endReference.blockId === undefined) { + throw new Error("Failed to map boundary matches back to raw messages") + } + boundaryOptionalIds.add(endReference.blockId) + } + + const strictRequiredIds = requiredBlockIds.filter((id) => !boundaryOptionalIds.has(id)) + const requiredSet = new Set(requiredBlockIds) + const keptPlaceholderIds = new Set() + const validPlaceholders: ParsedBlockPlaceholder[] = [] + + for (const placeholder of placeholders) { + const isKnown = summaryByBlockId.has(placeholder.blockId) + const isRequired = requiredSet.has(placeholder.blockId) + const isDuplicate = keptPlaceholderIds.has(placeholder.blockId) + + if (isKnown && isRequired && !isDuplicate) { + validPlaceholders.push(placeholder) + keptPlaceholderIds.add(placeholder.blockId) + } + } + + placeholders.length = 0 + placeholders.push(...validPlaceholders) + + return strictRequiredIds.filter((id) => !keptPlaceholderIds.has(id)) +} + +export function injectBlockPlaceholders( + summary: string, + placeholders: ParsedBlockPlaceholder[], + summaryByBlockId: Map, + startReference: BoundaryReference, + endReference: BoundaryReference, +): InjectedSummaryResult { + let cursor = 0 + let expanded = summary + const consumed: number[] = [] + const consumedSeen = new Set() + + if (placeholders.length > 0) { + expanded = "" + for (const placeholder of placeholders) { + const target = summaryByBlockId.get(placeholder.blockId) + if (!target) { + throw new Error(`Compressed block not found: (b${placeholder.blockId})`) + } + + expanded += summary.slice(cursor, placeholder.startIndex) + expanded += restoreSummary(target.summary) + cursor = placeholder.endIndex + + if (!consumedSeen.has(placeholder.blockId)) { + consumedSeen.add(placeholder.blockId) + consumed.push(placeholder.blockId) + } + } + + expanded += summary.slice(cursor) + } + + expanded = injectBoundarySummary( + expanded, + startReference, + "start", + summaryByBlockId, + consumed, + consumedSeen, + ) + expanded = injectBoundarySummary( + expanded, + endReference, + "end", + summaryByBlockId, + consumed, + consumedSeen, + ) + + return { + expandedSummary: expanded, + consumedBlockIds: consumed, + } +} + +export function appendMissingBlockSummaries( + summary: string, + missingBlockIds: number[], + summaryByBlockId: Map, + consumedBlockIds: number[], +): InjectedSummaryResult { + const consumedSeen = new Set(consumedBlockIds) + const consumed = [...consumedBlockIds] + + const missingSummaries: string[] = [] + for (const blockId of missingBlockIds) { + if (consumedSeen.has(blockId)) { + continue + } + + const target = summaryByBlockId.get(blockId) + if (!target) { + throw new Error(`Compressed block not found: (b${blockId})`) + } + + missingSummaries.push(`\n### (b${blockId})\n${restoreSummary(target.summary)}`) + consumedSeen.add(blockId) + consumed.push(blockId) + } + + if (missingSummaries.length === 0) { + return { + expandedSummary: summary, + consumedBlockIds: consumed, + } + } + + const heading = + "\n\nThe following previously compressed summaries were also part of this conversation section:" + + return { + expandedSummary: summary + heading + missingSummaries.join(""), + consumedBlockIds: consumed, + } +} + +function restoreSummary(summary: string): string { + const headerMatch = summary.match(/^\s*\[Compressed conversation(?: section)?(?: b\d+)?\]/i) + if (!headerMatch) { + return summary + } + + const afterHeader = summary.slice(headerMatch[0].length) + const withoutLeadingBreaks = afterHeader.replace(/^(?:\r?\n)+/, "") + return withoutLeadingBreaks + .replace(/(?:\r?\n)*b\d+<\/dcp-message-id>\s*$/i, "") + .replace(/(?:\r?\n)+$/, "") +} + +function injectBoundarySummary( + summary: string, + reference: BoundaryReference, + position: "start" | "end", + summaryByBlockId: Map, + consumed: number[], + consumedSeen: Set, +): string { + if (reference.kind !== "compressed-block" || reference.blockId === undefined) { + return summary + } + if (consumedSeen.has(reference.blockId)) { + return summary + } + + const target = summaryByBlockId.get(reference.blockId) + if (!target) { + throw new Error(`Compressed block not found: (b${reference.blockId})`) + } + + const injectedBody = restoreSummary(target.summary) + const left = position === "start" ? injectedBody.trim() : summary.trim() + const right = position === "start" ? summary.trim() : injectedBody.trim() + const next = !left ? right : !right ? left : `${left}\n\n${right}` + + consumedSeen.add(reference.blockId) + consumed.push(reference.blockId) + return next +} diff --git a/lib/compress/range.ts b/lib/compress/range.ts new file mode 100644 index 00000000..a7b544b0 --- /dev/null +++ b/lib/compress/range.ts @@ -0,0 +1,174 @@ +import { tool } from "@opencode-ai/plugin" +import type { ToolContext } from "./types" +import { countTokens } from "../strategies/utils" +import { RANGE_FORMAT_OVERLAY } from "../prompts/internal-overlays" +import { finalizeSession, prepareSession, type NotificationEntry } from "./pipeline" +import { appendProtectedTools, appendProtectedUserMessages } from "./protected-content" +import { + appendMissingBlockSummaries, + injectBlockPlaceholders, + parseBlockPlaceholders, + resolveRanges, + validateArgs, + validateNonOverlapping, + validateSummaryPlaceholders, +} from "./range-utils" +import { + COMPRESSED_BLOCK_HEADER, + allocateBlockId, + allocateRunId, + applyCompressionState, + wrapCompressedSummary, +} from "./state" +import type { CompressRangeToolArgs } from "./types" + +function buildSchema() { + return { + topic: tool.schema + .string() + .describe("Short label (3-5 words) for display - e.g., 'Auth System Exploration'"), + content: tool.schema + .array( + tool.schema.object({ + startId: tool.schema + .string() + .describe( + "Message or block ID marking the beginning of range (e.g. m0001, b2)", + ), + endId: tool.schema + .string() + .describe("Message or block ID marking the end of range (e.g. m0012, b5)"), + summary: tool.schema + .string() + .describe("Complete technical summary replacing all content in range"), + }), + ) + .describe( + "One or more ranges to compress, each with start/end boundaries and a summary", + ), + } +} + +export function createCompressRangeTool(ctx: ToolContext): ReturnType { + ctx.prompts.reload() + const runtimePrompts = ctx.prompts.getRuntimePrompts() + + return tool({ + description: runtimePrompts.compressRange + RANGE_FORMAT_OVERLAY, + args: buildSchema(), + async execute(args, toolCtx) { + const input = args as CompressRangeToolArgs + validateArgs(input) + + const { rawMessages, searchContext } = await prepareSession( + ctx, + toolCtx, + `Compress Range: ${input.topic}`, + ) + const resolvedPlans = resolveRanges(input, searchContext, ctx.state) + validateNonOverlapping(resolvedPlans) + + const notifications: NotificationEntry[] = [] + const preparedPlans: Array<{ + entry: (typeof resolvedPlans)[number]["entry"] + selection: (typeof resolvedPlans)[number]["selection"] + anchorMessageId: string + finalSummary: string + consumedBlockIds: number[] + }> = [] + let totalCompressedMessages = 0 + + for (const plan of resolvedPlans) { + const parsedPlaceholders = parseBlockPlaceholders(plan.entry.summary) + const missingBlockIds = validateSummaryPlaceholders( + parsedPlaceholders, + plan.selection.requiredBlockIds, + plan.selection.startReference, + plan.selection.endReference, + searchContext.summaryByBlockId, + ) + + const injected = injectBlockPlaceholders( + plan.entry.summary, + parsedPlaceholders, + searchContext.summaryByBlockId, + plan.selection.startReference, + plan.selection.endReference, + ) + + const summaryWithUsers = appendProtectedUserMessages( + injected.expandedSummary, + plan.selection, + searchContext, + ctx.state, + ctx.config.compress.protectUserMessages, + ) + + const summaryWithTools = await appendProtectedTools( + ctx.client, + ctx.state, + ctx.config.experimental.allowSubAgents, + summaryWithUsers, + plan.selection, + searchContext, + ctx.config.compress.protectedTools, + ctx.config.protectedFilePatterns, + ) + + const completedSummary = appendMissingBlockSummaries( + summaryWithTools, + missingBlockIds, + searchContext.summaryByBlockId, + injected.consumedBlockIds, + ) + + preparedPlans.push({ + entry: plan.entry, + selection: plan.selection, + anchorMessageId: plan.anchorMessageId, + finalSummary: completedSummary.expandedSummary, + consumedBlockIds: completedSummary.consumedBlockIds, + }) + } + + const runId = allocateRunId(ctx.state) + + for (const preparedPlan of preparedPlans) { + const blockId = allocateBlockId(ctx.state) + const storedSummary = wrapCompressedSummary(blockId, preparedPlan.finalSummary) + const summaryTokens = countTokens(storedSummary) + + const applied = applyCompressionState( + ctx.state, + { + topic: input.topic, + batchTopic: input.topic, + startId: preparedPlan.entry.startId, + endId: preparedPlan.entry.endId, + mode: "range", + runId, + compressMessageId: toolCtx.messageID, + }, + preparedPlan.selection, + preparedPlan.anchorMessageId, + blockId, + storedSummary, + preparedPlan.consumedBlockIds, + ) + + totalCompressedMessages += applied.messageIds.length + + notifications.push({ + blockId, + runId, + summary: preparedPlan.finalSummary, + summaryTokens, + }) + } + + await finalizeSession(ctx, toolCtx, rawMessages, notifications, input.topic) + + return `Compressed ${totalCompressedMessages} messages into ${COMPRESSED_BLOCK_HEADER}.` + }, + }) +} diff --git a/lib/compress/search.ts b/lib/compress/search.ts new file mode 100644 index 00000000..c53ffb40 --- /dev/null +++ b/lib/compress/search.ts @@ -0,0 +1,267 @@ +import type { SessionState, WithParts } from "../state" +import { formatBlockRef, parseBoundaryId } from "../message-ids" +import { isIgnoredUserMessage } from "../messages/utils" +import { countAllMessageTokens } from "../strategies/utils" +import type { BoundaryReference, SearchContext, SelectionResolution } from "./types" + +export async function fetchSessionMessages(client: any, sessionId: string): Promise { + const response = await client.session.messages({ + path: { id: sessionId }, + }) + + const payload = (response?.data || response) as WithParts[] + return Array.isArray(payload) ? payload : [] +} + +export function buildSearchContext(state: SessionState, rawMessages: WithParts[]): SearchContext { + const rawMessagesById = new Map() + const rawIndexById = new Map() + for (const msg of rawMessages) { + rawMessagesById.set(msg.info.id, msg) + } + for (let index = 0; index < rawMessages.length; index++) { + const message = rawMessages[index] + if (!message) { + continue + } + rawIndexById.set(message.info.id, index) + } + + const summaryByBlockId = new Map() + for (const [blockId, block] of state.prune.messages.blocksById) { + if (!block.active) { + continue + } + summaryByBlockId.set(blockId, block) + } + + return { + rawMessages, + rawMessagesById, + rawIndexById, + summaryByBlockId, + } +} + +export function resolveBoundaryIds( + context: SearchContext, + state: SessionState, + startId: string, + endId: string, +): { startReference: BoundaryReference; endReference: BoundaryReference } { + const lookup = buildBoundaryLookup(context, state) + const issues: string[] = [] + const parsedStartId = parseBoundaryId(startId) + const parsedEndId = parseBoundaryId(endId) + + if (parsedStartId === null) { + issues.push("startId is invalid. Use an injected message ID (mNNNN) or block ID (bN).") + } + + if (parsedEndId === null) { + issues.push("endId is invalid. Use an injected message ID (mNNNN) or block ID (bN).") + } + + if (issues.length > 0) { + throw new Error( + issues.length === 1 ? issues[0] : issues.map((issue) => `- ${issue}`).join("\n"), + ) + } + + if (!parsedStartId || !parsedEndId) { + throw new Error("Invalid boundary ID(s)") + } + + const startReference = lookup.get(parsedStartId.ref) + const endReference = lookup.get(parsedEndId.ref) + + if (!startReference) { + issues.push( + `startId ${parsedStartId.ref} is not available in the current conversation context. Choose an injected ID visible in context.`, + ) + } + + if (!endReference) { + issues.push( + `endId ${parsedEndId.ref} is not available in the current conversation context. Choose an injected ID visible in context.`, + ) + } + + if (issues.length > 0) { + throw new Error( + issues.length === 1 ? issues[0] : issues.map((issue) => `- ${issue}`).join("\n"), + ) + } + + if (!startReference || !endReference) { + throw new Error("Failed to resolve boundary IDs") + } + + if (startReference.rawIndex > endReference.rawIndex) { + throw new Error( + `startId ${parsedStartId.ref} appears after endId ${parsedEndId.ref} in the conversation. Start must come before end.`, + ) + } + + return { startReference, endReference } +} + +export function resolveSelection( + context: SearchContext, + startReference: BoundaryReference, + endReference: BoundaryReference, +): SelectionResolution { + const startRawIndex = startReference.rawIndex + const endRawIndex = endReference.rawIndex + const messageIds: string[] = [] + const messageSeen = new Set() + const toolIds: string[] = [] + const toolSeen = new Set() + const requiredBlockIds: number[] = [] + const requiredBlockSeen = new Set() + const messageTokenById = new Map() + + for (let index = startRawIndex; index <= endRawIndex; index++) { + const rawMessage = context.rawMessages[index] + if (!rawMessage) { + continue + } + if (rawMessage.info.role === "user" && isIgnoredUserMessage(rawMessage)) { + continue + } + + const messageId = rawMessage.info.id + if (!messageSeen.has(messageId)) { + messageSeen.add(messageId) + messageIds.push(messageId) + } + + if (!messageTokenById.has(messageId)) { + messageTokenById.set(messageId, countAllMessageTokens(rawMessage)) + } + + const parts = Array.isArray(rawMessage.parts) ? rawMessage.parts : [] + for (const part of parts) { + if (part.type !== "tool" || !part.callID) { + continue + } + if (toolSeen.has(part.callID)) { + continue + } + toolSeen.add(part.callID) + toolIds.push(part.callID) + } + } + + const selectedMessageIds = new Set(messageIds) + const summariesInSelection: Array<{ blockId: number; rawIndex: number }> = [] + for (const summary of context.summaryByBlockId.values()) { + if (!selectedMessageIds.has(summary.anchorMessageId)) { + continue + } + + const anchorIndex = context.rawIndexById.get(summary.anchorMessageId) + if (anchorIndex === undefined) { + continue + } + + summariesInSelection.push({ + blockId: summary.blockId, + rawIndex: anchorIndex, + }) + } + + summariesInSelection.sort((a, b) => a.rawIndex - b.rawIndex || a.blockId - b.blockId) + for (const summary of summariesInSelection) { + if (requiredBlockSeen.has(summary.blockId)) { + continue + } + requiredBlockSeen.add(summary.blockId) + requiredBlockIds.push(summary.blockId) + } + + if (messageIds.length === 0) { + throw new Error( + "Failed to map boundary matches back to raw messages. Choose boundaries that include original conversation messages.", + ) + } + + return { + startReference, + endReference, + messageIds, + messageTokenById, + toolIds, + requiredBlockIds, + } +} + +export function resolveAnchorMessageId(startReference: BoundaryReference): string { + if (startReference.kind === "compressed-block") { + if (!startReference.anchorMessageId) { + throw new Error("Failed to map boundary matches back to raw messages") + } + return startReference.anchorMessageId + } + + if (!startReference.messageId) { + throw new Error("Failed to map boundary matches back to raw messages") + } + return startReference.messageId +} + +function buildBoundaryLookup( + context: SearchContext, + state: SessionState, +): Map { + const lookup = new Map() + + for (const [messageRef, messageId] of state.messageIds.byRef) { + const rawMessage = context.rawMessagesById.get(messageId) + if (!rawMessage) { + continue + } + if (rawMessage.info.role === "user" && isIgnoredUserMessage(rawMessage)) { + continue + } + + const rawIndex = context.rawIndexById.get(messageId) + if (rawIndex === undefined) { + continue + } + lookup.set(messageRef, { + kind: "message", + rawIndex, + messageId, + }) + } + + const summaries = Array.from(context.summaryByBlockId.values()).sort( + (a, b) => a.blockId - b.blockId, + ) + for (const summary of summaries) { + const anchorMessage = context.rawMessagesById.get(summary.anchorMessageId) + if (!anchorMessage) { + continue + } + if (anchorMessage.info.role === "user" && isIgnoredUserMessage(anchorMessage)) { + continue + } + + const rawIndex = context.rawIndexById.get(summary.anchorMessageId) + if (rawIndex === undefined) { + continue + } + const blockRef = formatBlockRef(summary.blockId) + if (!lookup.has(blockRef)) { + lookup.set(blockRef, { + kind: "compressed-block", + rawIndex, + blockId: summary.blockId, + anchorMessageId: summary.anchorMessageId, + }) + } + } + + return lookup +} diff --git a/lib/compress/state.ts b/lib/compress/state.ts new file mode 100644 index 00000000..ca6df410 --- /dev/null +++ b/lib/compress/state.ts @@ -0,0 +1,242 @@ +import type { CompressionBlock, SessionState } from "../state" +import { formatBlockRef, formatMessageIdTag } from "../message-ids" +import type { AppliedCompressionResult, CompressionStateInput, SelectionResolution } from "./types" + +export const COMPRESSED_BLOCK_HEADER = "[Compressed conversation section]" + +export function allocateBlockId(state: SessionState): number { + const next = state.prune.messages.nextBlockId + if (!Number.isInteger(next) || next < 1) { + state.prune.messages.nextBlockId = 2 + return 1 + } + + state.prune.messages.nextBlockId = next + 1 + return next +} + +export function allocateRunId(state: SessionState): number { + const next = state.prune.messages.nextRunId + if (!Number.isInteger(next) || next < 1) { + state.prune.messages.nextRunId = 2 + return 1 + } + + state.prune.messages.nextRunId = next + 1 + return next +} + +export function wrapCompressedSummary(blockId: number, summary: string): string { + const header = COMPRESSED_BLOCK_HEADER + const footer = formatMessageIdTag(formatBlockRef(blockId)) + const body = summary.trim() + if (body.length === 0) { + return `${header}\n${footer}` + } + return `${header}\n${body}\n\n${footer}` +} + +export function applyCompressionState( + state: SessionState, + input: CompressionStateInput, + selection: SelectionResolution, + anchorMessageId: string, + blockId: number, + summary: string, + consumedBlockIds: number[], +): AppliedCompressionResult { + const messagesState = state.prune.messages + const consumed = [...new Set(consumedBlockIds.filter((id) => Number.isInteger(id) && id > 0))] + const included = [...consumed] + + const effectiveMessageIds = new Set(selection.messageIds) + const effectiveToolIds = new Set(selection.toolIds) + + for (const consumedBlockId of consumed) { + const consumedBlock = messagesState.blocksById.get(consumedBlockId) + if (!consumedBlock) { + continue + } + for (const messageId of consumedBlock.effectiveMessageIds) { + effectiveMessageIds.add(messageId) + } + for (const toolId of consumedBlock.effectiveToolIds) { + effectiveToolIds.add(toolId) + } + } + + const initiallyActiveMessages = new Set() + for (const messageId of effectiveMessageIds) { + const entry = messagesState.byMessageId.get(messageId) + if (entry && entry.activeBlockIds.length > 0) { + initiallyActiveMessages.add(messageId) + } + } + + const initiallyActiveToolIds = new Set() + for (const activeBlockId of messagesState.activeBlockIds) { + const activeBlock = messagesState.blocksById.get(activeBlockId) + if (!activeBlock || !activeBlock.active) { + continue + } + + for (const toolId of activeBlock.effectiveToolIds) { + initiallyActiveToolIds.add(toolId) + } + } + + const createdAt = Date.now() + const block: CompressionBlock = { + blockId, + runId: input.runId, + active: true, + deactivatedByUser: false, + compressedTokens: 0, + mode: input.mode, + topic: input.topic, + batchTopic: input.batchTopic, + startId: input.startId, + endId: input.endId, + anchorMessageId, + compressMessageId: input.compressMessageId, + includedBlockIds: included, + consumedBlockIds: consumed, + parentBlockIds: [], + directMessageIds: [], + directToolIds: [], + effectiveMessageIds: [...effectiveMessageIds], + effectiveToolIds: [...effectiveToolIds], + createdAt, + summary, + } + + messagesState.blocksById.set(blockId, block) + messagesState.activeBlockIds.add(blockId) + messagesState.activeByAnchorMessageId.set(anchorMessageId, blockId) + + const deactivatedAt = Date.now() + for (const consumedBlockId of consumed) { + const consumedBlock = messagesState.blocksById.get(consumedBlockId) + if (!consumedBlock || !consumedBlock.active) { + continue + } + + consumedBlock.active = false + consumedBlock.deactivatedAt = deactivatedAt + consumedBlock.deactivatedByBlockId = blockId + if (!consumedBlock.parentBlockIds.includes(blockId)) { + consumedBlock.parentBlockIds.push(blockId) + } + + messagesState.activeBlockIds.delete(consumedBlockId) + const mappedBlockId = messagesState.activeByAnchorMessageId.get( + consumedBlock.anchorMessageId, + ) + if (mappedBlockId === consumedBlockId) { + messagesState.activeByAnchorMessageId.delete(consumedBlock.anchorMessageId) + } + } + + const removeActiveBlockId = ( + entry: { activeBlockIds: number[] }, + blockIdToRemove: number, + ): void => { + if (entry.activeBlockIds.length === 0) { + return + } + entry.activeBlockIds = entry.activeBlockIds.filter((id) => id !== blockIdToRemove) + } + + for (const consumedBlockId of consumed) { + const consumedBlock = messagesState.blocksById.get(consumedBlockId) + if (!consumedBlock) { + continue + } + for (const messageId of consumedBlock.effectiveMessageIds) { + const entry = messagesState.byMessageId.get(messageId) + if (!entry) { + continue + } + removeActiveBlockId(entry, consumedBlockId) + } + } + + for (const messageId of selection.messageIds) { + const tokenCount = selection.messageTokenById.get(messageId) || 0 + const existing = messagesState.byMessageId.get(messageId) + + if (!existing) { + messagesState.byMessageId.set(messageId, { + tokenCount, + allBlockIds: [blockId], + activeBlockIds: [blockId], + }) + continue + } + + existing.tokenCount = Math.max(existing.tokenCount, tokenCount) + if (!existing.allBlockIds.includes(blockId)) { + existing.allBlockIds.push(blockId) + } + if (!existing.activeBlockIds.includes(blockId)) { + existing.activeBlockIds.push(blockId) + } + } + + for (const messageId of block.effectiveMessageIds) { + if (selection.messageTokenById.has(messageId)) { + continue + } + + const existing = messagesState.byMessageId.get(messageId) + if (!existing) { + continue + } + if (!existing.allBlockIds.includes(blockId)) { + existing.allBlockIds.push(blockId) + } + if (!existing.activeBlockIds.includes(blockId)) { + existing.activeBlockIds.push(blockId) + } + } + + let compressedTokens = 0 + const newlyCompressedMessageIds: string[] = [] + for (const messageId of effectiveMessageIds) { + const entry = messagesState.byMessageId.get(messageId) + if (!entry) { + continue + } + + const isNowActive = entry.activeBlockIds.length > 0 + const wasActive = initiallyActiveMessages.has(messageId) + + if (isNowActive && !wasActive) { + compressedTokens += entry.tokenCount + newlyCompressedMessageIds.push(messageId) + } + } + + const newlyCompressedToolIds: string[] = [] + for (const toolId of effectiveToolIds) { + if (!initiallyActiveToolIds.has(toolId)) { + newlyCompressedToolIds.push(toolId) + } + } + + block.directMessageIds = [...newlyCompressedMessageIds] + block.directToolIds = [...newlyCompressedToolIds] + + block.compressedTokens = compressedTokens + + state.stats.pruneTokenCounter += compressedTokens + state.stats.totalPruneTokens += state.stats.pruneTokenCounter + state.stats.pruneTokenCounter = 0 + + return { + compressedTokens, + messageIds: selection.messageIds, + newlyCompressedMessageIds, + newlyCompressedToolIds, + } +} diff --git a/lib/compress/types.ts b/lib/compress/types.ts new file mode 100644 index 00000000..dc839c45 --- /dev/null +++ b/lib/compress/types.ts @@ -0,0 +1,105 @@ +import type { PluginConfig } from "../config" +import type { Logger } from "../logger" +import type { PromptStore } from "../prompts/store" +import type { CompressionBlock, CompressionMode, SessionState, WithParts } from "../state" + +export interface ToolContext { + client: any + state: SessionState + logger: Logger + config: PluginConfig + prompts: PromptStore +} + +export interface CompressRangeEntry { + startId: string + endId: string + summary: string +} + +export interface CompressRangeToolArgs { + topic: string + content: CompressRangeEntry[] +} + +export interface CompressMessageEntry { + messageId: string + topic: string + summary: string +} + +export interface CompressMessageToolArgs { + topic: string + content: CompressMessageEntry[] +} + +export interface BoundaryReference { + kind: "message" | "compressed-block" + rawIndex: number + messageId?: string + blockId?: number + anchorMessageId?: string +} + +export interface SearchContext { + rawMessages: WithParts[] + rawMessagesById: Map + rawIndexById: Map + summaryByBlockId: Map +} + +export interface SelectionResolution { + startReference: BoundaryReference + endReference: BoundaryReference + messageIds: string[] + messageTokenById: Map + toolIds: string[] + requiredBlockIds: number[] +} + +export interface ResolvedMessageCompression { + entry: CompressMessageEntry + selection: SelectionResolution + anchorMessageId: string +} + +export interface ResolvedRangeCompression { + index: number + entry: CompressRangeEntry + selection: SelectionResolution + anchorMessageId: string +} + +export interface ResolvedMessageCompressionsResult { + plans: ResolvedMessageCompression[] + skippedIssues: string[] +} + +export interface ParsedBlockPlaceholder { + raw: string + blockId: number + startIndex: number + endIndex: number +} + +export interface InjectedSummaryResult { + expandedSummary: string + consumedBlockIds: number[] +} + +export interface AppliedCompressionResult { + compressedTokens: number + messageIds: string[] + newlyCompressedMessageIds: string[] + newlyCompressedToolIds: string[] +} + +export interface CompressionStateInput { + topic: string + batchTopic: string + startId: string + endId: string + mode: CompressionMode + runId: number + compressMessageId: string +} diff --git a/lib/config.ts b/lib/config.ts index c0c96547..f59fd358 100644 --- a/lib/config.ts +++ b/lib/config.ts @@ -5,13 +5,15 @@ import { parse } from "jsonc-parser" import type { PluginInput } from "@opencode-ai/plugin" type Permission = "ask" | "allow" | "deny" +type CompressMode = "range" | "message" export interface Deduplication { enabled: boolean protectedTools: string[] } -export interface CompressTool { +export interface CompressConfig { + mode: CompressMode permission: Permission showCompression: boolean maxContextLimit: number | `${number}%` @@ -21,7 +23,6 @@ export interface CompressTool { nudgeFrequency: number iterationNudgeThreshold: number nudgeForce: "strong" | "soft" - flatSchema: boolean protectedTools: string[] protectUserMessages: boolean } @@ -36,10 +37,6 @@ export interface ManualModeConfig { automaticStrategies: boolean } -export interface SupersedeWrites { - enabled: boolean -} - export interface PurgeErrors { enabled: boolean turns: number @@ -66,15 +63,14 @@ export interface PluginConfig { turnProtection: TurnProtection experimental: ExperimentalConfig protectedFilePatterns: string[] - compress: CompressTool + compress: CompressConfig strategies: { deduplication: Deduplication - supersedeWrites: SupersedeWrites purgeErrors: PurgeErrors } } -type CompressOverride = Partial +type CompressOverride = Partial const DEFAULT_PROTECTED_TOOLS = [ "task", @@ -112,6 +108,7 @@ export const VALID_CONFIG_KEYS = new Set([ "manualMode.enabled", "manualMode.automaticStrategies", "compress", + "compress.mode", "compress.permission", "compress.showCompression", "compress.maxContextLimit", @@ -121,15 +118,12 @@ export const VALID_CONFIG_KEYS = new Set([ "compress.nudgeFrequency", "compress.iterationNudgeThreshold", "compress.nudgeForce", - "compress.flatSchema", "compress.protectedTools", "compress.protectUserMessages", "strategies", "strategies.deduplication", "strategies.deduplication.enabled", "strategies.deduplication.protectedTools", - "strategies.supersedeWrites", - "strategies.supersedeWrites.enabled", "strategies.purgeErrors", "strategies.purgeErrors.enabled", "strategies.purgeErrors.turns", @@ -347,6 +341,18 @@ export function validateConfigTypes(config: Record): ValidationErro actual: typeof compress, }) } else { + if ( + compress.mode !== undefined && + compress.mode !== "range" && + compress.mode !== "message" + ) { + errors.push({ + key: "compress.mode", + expected: '"range" | "message"', + actual: JSON.stringify(compress.mode), + }) + } + if ( compress.nudgeFrequency !== undefined && typeof compress.nudgeFrequency !== "number" @@ -389,14 +395,6 @@ export function validateConfigTypes(config: Record): ValidationErro }) } - if (compress.flatSchema !== undefined && typeof compress.flatSchema !== "boolean") { - errors.push({ - key: "compress.flatSchema", - expected: "boolean", - actual: typeof compress.flatSchema, - }) - } - if (compress.protectedTools !== undefined && !Array.isArray(compress.protectedTools)) { errors.push({ key: "compress.protectedTools", @@ -532,19 +530,6 @@ export function validateConfigTypes(config: Record): ValidationErro }) } - if (strategies.supersedeWrites) { - if ( - strategies.supersedeWrites.enabled !== undefined && - typeof strategies.supersedeWrites.enabled !== "boolean" - ) { - errors.push({ - key: "strategies.supersedeWrites.enabled", - expected: "boolean", - actual: typeof strategies.supersedeWrites.enabled, - }) - } - } - if (strategies.purgeErrors) { if ( strategies.purgeErrors.enabled !== undefined && @@ -662,6 +647,7 @@ const defaultConfig: PluginConfig = { }, protectedFilePatterns: [], compress: { + mode: "range", permission: "allow", showCompression: false, maxContextLimit: 150000, @@ -669,7 +655,6 @@ const defaultConfig: PluginConfig = { nudgeFrequency: 5, iterationNudgeThreshold: 15, nudgeForce: "soft", - flatSchema: false, protectedTools: [...COMPRESS_DEFAULT_PROTECTED_TOOLS], protectUserMessages: false, }, @@ -678,9 +663,6 @@ const defaultConfig: PluginConfig = { enabled: true, protectedTools: [], }, - supersedeWrites: { - enabled: true, - }, purgeErrors: { enabled: true, turns: 4, @@ -805,9 +787,6 @@ function mergeStrategies( ]), ], }, - supersedeWrites: { - enabled: override.supersedeWrites?.enabled ?? base.supersedeWrites.enabled, - }, purgeErrors: { enabled: override.purgeErrors?.enabled ?? base.purgeErrors.enabled, turns: override.purgeErrors?.turns ?? base.purgeErrors.turns, @@ -830,6 +809,7 @@ function mergeCompress( } return { + mode: override.mode ?? base.mode, permission: override.permission ?? base.permission, showCompression: override.showCompression ?? base.showCompression, maxContextLimit: override.maxContextLimit ?? base.maxContextLimit, @@ -839,7 +819,6 @@ function mergeCompress( nudgeFrequency: override.nudgeFrequency ?? base.nudgeFrequency, iterationNudgeThreshold: override.iterationNudgeThreshold ?? base.iterationNudgeThreshold, nudgeForce: override.nudgeForce ?? base.nudgeForce, - flatSchema: override.flatSchema ?? base.flatSchema, protectedTools: [...new Set([...base.protectedTools, ...(override.protectedTools ?? [])])], protectUserMessages: override.protectUserMessages ?? base.protectUserMessages, } @@ -908,7 +887,6 @@ function deepCloneConfig(config: PluginConfig): PluginConfig { ...config.strategies.deduplication, protectedTools: [...config.strategies.deduplication.protectedTools], }, - supersedeWrites: { ...config.strategies.supersedeWrites }, purgeErrors: { ...config.strategies.purgeErrors, protectedTools: [...config.strategies.purgeErrors.protectedTools], diff --git a/lib/hooks.ts b/lib/hooks.ts index 3fcc5e63..8b3d9f1d 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -2,6 +2,7 @@ import type { SessionState, WithParts } from "./state" import type { Logger } from "./logger" import type { PluginConfig } from "./config" import { assignMessageRefs } from "./message-ids" +import { buildPriorityMap } from "./messages/priority" import { syncToolCache } from "./state/tool-cache" import { prune, @@ -11,7 +12,12 @@ import { injectExtendedSubAgentResults, stripStaleMetadata, } from "./messages" -import { buildToolIdList, isIgnoredUserMessage, stripHallucinations } from "./messages/utils" +import { + buildToolIdList, + isIgnoredUserMessage, + stripHallucinations, + stripHallucinationsFromString, +} from "./messages/utils" import { checkSession } from "./state" import { renderSystemPrompt } from "./prompts" import { handleStatsCommand } from "./commands/stats" @@ -33,9 +39,6 @@ const INTERNAL_AGENT_SIGNATURES = [ "Summarize what was done in this conversation", ] -const DCP_MESSAGE_ID_TAG_REGEX = /(?:m\d+|b\d+)<\/dcp-message-id>/g -const DCP_SYSTEM_REMINDER_REGEX = /]*>[\s\S]*?<\/dcp-system-reminder>/g - function applyManualPrompt(state: SessionState, messages: WithParts[], logger: Logger): void { const pending = state.pendingManualTrigger if (!pending) { @@ -148,9 +151,17 @@ export function createChatMessageTransformHandler( output.messages, config.experimental.allowSubAgents, ) + const compressionPriorities = buildPriorityMap(config, state, output.messages) prompts.reload() - injectCompressNudges(state, config, logger, output.messages, prompts.getRuntimePrompts()) - injectMessageIds(state, config, output.messages) + injectCompressNudges( + state, + config, + logger, + output.messages, + prompts.getRuntimePrompts(), + compressionPriorities, + ) + injectMessageIds(state, config, output.messages, compressionPriorities) applyManualPrompt(state, output.messages, logger) stripStaleMetadata(output.messages) @@ -280,8 +291,6 @@ export function createTextCompleteHandler() { _input: { sessionID: string; messageID: string; partID: string }, output: { text: string }, ) => { - output.text = output.text - .replace(DCP_SYSTEM_REMINDER_REGEX, "") - .replace(DCP_MESSAGE_ID_TAG_REGEX, "") + output.text = stripHallucinationsFromString(output.text) } } diff --git a/lib/message-ids.ts b/lib/message-ids.ts index 9f5b60e6..edb04cca 100644 --- a/lib/message-ids.ts +++ b/lib/message-ids.ts @@ -90,8 +90,30 @@ export function parseBoundaryId(id: string): ParsedBoundaryId | null { return null } -export function formatMessageIdTag(ref: string): string { - return `\n<${MESSAGE_ID_TAG_NAME}>${ref}` +function escapeXmlAttribute(value: string): string { + return value + .replace(/&/g, "&") + .replace(/"/g, """) + .replace(//g, ">") +} + +export function formatMessageIdTag( + ref: string, + attributes?: Record, +): string { + const serializedAttributes = Object.entries(attributes || {}) + .sort(([left], [right]) => left.localeCompare(right)) + .map(([name, value]) => { + if (name.trim().length === 0 || typeof value !== "string" || value.length === 0) { + return "" + } + + return ` ${name}="${escapeXmlAttribute(value)}"` + }) + .join("") + + return `\n<${MESSAGE_ID_TAG_NAME}${serializedAttributes}>${ref}` } export function assignMessageRefs(state: SessionState, messages: WithParts[]): number { diff --git a/lib/messages/inject/inject.ts b/lib/messages/inject/inject.ts index 17a4286a..67101716 100644 --- a/lib/messages/inject/inject.ts +++ b/lib/messages/inject/inject.ts @@ -3,13 +3,16 @@ import type { Logger } from "../../logger" import type { PluginConfig } from "../../config" import type { RuntimePrompts } from "../../prompts/store" import { formatMessageIdTag } from "../../message-ids" +import type { CompressionPriorityMap } from "../priority" import { compressPermission, getLastUserMessage } from "../../shared-utils" import { saveSessionState } from "../../state/persistence" import { + appendToTextPart, appendIdToTool, createSyntheticTextPart, findLastToolPart, isIgnoredUserMessage, + isProtectedUserMessage, } from "../utils" import { addAnchor, @@ -29,6 +32,7 @@ export const injectCompressNudges = ( logger: Logger, messages: WithParts[], prompts: RuntimePrompts, + compressionPriorities?: CompressionPriorityMap, ): void => { if (compressPermission(state, config) === "deny") { return @@ -127,7 +131,7 @@ export const injectCompressNudges = ( } } - applyAnchoredNudges(state, config, messages, prompts) + applyAnchoredNudges(state, config, messages, prompts, compressionPriorities) if (anchorsChanged) { void saveSessionState(state, logger) @@ -138,6 +142,7 @@ export const injectMessageIds = ( state: SessionState, config: PluginConfig, messages: WithParts[], + compressionPriorities?: CompressionPriorityMap, ): void => { if (compressPermission(state, config) === "deny") { return @@ -153,9 +158,21 @@ export const injectMessageIds = ( continue } - const tag = formatMessageIdTag(messageRef) + const isBlockedMessage = isProtectedUserMessage(config, message) + const priority = + config.compress.mode === "message" && !isBlockedMessage + ? compressionPriorities?.get(message.info.id)?.priority + : undefined + const tag = formatMessageIdTag( + isBlockedMessage ? "BLOCKED" : messageRef, + priority ? { priority } : undefined, + ) if (message.info.role === "user") { + if (appendToTextPart(message, tag)) { + continue + } + message.parts.push(createSyntheticTextPart(message, tag)) continue } @@ -169,6 +186,10 @@ export const injectMessageIds = ( continue } + if (appendToTextPart(message, tag)) { + continue + } + const syntheticPart = createSyntheticTextPart(message, tag) const firstToolIndex = message.parts.findIndex((p) => p.type === "tool") if (firstToolIndex === -1) { diff --git a/lib/messages/inject/utils.ts b/lib/messages/inject/utils.ts index bdf939e8..72adfedd 100644 --- a/lib/messages/inject/utils.ts +++ b/lib/messages/inject/utils.ts @@ -1,11 +1,19 @@ import type { SessionState, WithParts } from "../../state" import type { PluginConfig } from "../../config" +import { renderMessagePriorityGuidance } from "../../prompts/message-priority-guidance" import type { RuntimePrompts } from "../../prompts/store" import type { UserMessage } from "@opencode-ai/sdk/v2" -import { createSyntheticTextPart, isIgnoredUserMessage } from "../utils" +import { + type CompressionPriorityMap, + type MessagePriority, + listPriorityRefsBeforeIndex, +} from "../priority" +import { appendToTextPart, createSyntheticTextPart, isIgnoredUserMessage } from "../utils" import { getLastUserMessage } from "../../shared-utils" import { getCurrentTokenUsage } from "../../strategies/utils" +const MESSAGE_MODE_NUDGE_PRIORITY: MessagePriority = "high" + export interface LastUserModelContext { providerId: string | undefined modelId: string | undefined @@ -184,72 +192,94 @@ export function buildCompressedBlockGuidance(state: SessionState): string { ].join("\n") } -function appendGuidanceToDcpTag(hintText: string, guidance: string): string { +function appendGuidanceToDcpTag(nudgeText: string, guidance: string): string { + if (!guidance.trim()) { + return nudgeText + } + const closeTag = "" - const closeTagIndex = hintText.lastIndexOf(closeTag) + const closeTagIndex = nudgeText.lastIndexOf(closeTag) if (closeTagIndex === -1) { - return hintText + return nudgeText } - const beforeClose = hintText.slice(0, closeTagIndex).trimEnd() - const afterClose = hintText.slice(closeTagIndex) + const beforeClose = nudgeText.slice(0, closeTagIndex).trimEnd() + const afterClose = nudgeText.slice(closeTagIndex) return `${beforeClose}\n\n${guidance}\n${afterClose}` } -function applyAnchoredNudge( - anchorMessageIds: Set, +function buildMessagePriorityGuidance( messages: WithParts[], - hintText: string, -): void { - if (anchorMessageIds.size === 0) { + compressionPriorities: CompressionPriorityMap | undefined, + anchorIndex: number, + priority: MessagePriority, +): string { + if (!compressionPriorities || compressionPriorities.size === 0) { + return "" + } + + const refs = listPriorityRefsBeforeIndex(messages, compressionPriorities, anchorIndex, priority) + const priorityLabel = `${priority[0].toUpperCase()}${priority.slice(1)}` + + return renderMessagePriorityGuidance(priorityLabel, refs) +} + +function injectAnchoredNudge(message: WithParts, nudgeText: string): void { + if (!nudgeText.trim()) { return } - for (const anchorMessageId of anchorMessageIds) { - const messageIndex = messages.findIndex((message) => message.info.id === anchorMessageId) - if (messageIndex === -1) { - continue - } + if (appendToTextPart(message, nudgeText)) { + return + } - const message = messages[messageIndex] - if (message.info.role === "user") { - message.parts.push(createSyntheticTextPart(message, hintText)) - continue - } + if (message.info.role === "user") { + message.parts.push(createSyntheticTextPart(message, nudgeText)) + return + } - if (message.info.role !== "assistant") { + if (message.info.role !== "assistant") { + return + } + + const syntheticPart = createSyntheticTextPart(message, nudgeText) + const firstToolIndex = message.parts.findIndex((p) => p.type === "tool") + if (firstToolIndex === -1) { + message.parts.push(syntheticPart) + } else { + message.parts.splice(firstToolIndex, 0, syntheticPart) + } +} + +function collectAnchoredMessages( + anchorMessageIds: Set, + messages: WithParts[], +): Array<{ message: WithParts; index: number }> { + const anchoredMessages: Array<{ message: WithParts; index: number }> = [] + + for (const anchorMessageId of anchorMessageIds) { + const index = messages.findIndex((message) => message.info.id === anchorMessageId) + if (index === -1) { continue } - const syntheticPart = createSyntheticTextPart(message, hintText) - const firstToolIndex = message.parts.findIndex((p) => p.type === "tool") - if (firstToolIndex === -1) { - message.parts.push(syntheticPart) - } else { - message.parts.splice(firstToolIndex, 0, syntheticPart) - } + anchoredMessages.push({ + message: messages[index], + index, + }) } + + return anchoredMessages } -export function applyAnchoredNudges( +function collectTurnNudgeAnchors( state: SessionState, config: PluginConfig, messages: WithParts[], - prompts: RuntimePrompts, -): void { - const compressedBlockGuidance = buildCompressedBlockGuidance(state) - - const contextLimitNudge = appendGuidanceToDcpTag( - prompts.contextLimitNudge, - compressedBlockGuidance, - ) - - applyAnchoredNudge(state.nudges.contextLimitAnchors, messages, contextLimitNudge) - +): Set { const turnNudgeAnchors = new Set() const targetRole = config.compress.nudgeForce === "strong" ? "user" : "assistant" - const turnNudge = appendGuidanceToDcpTag(prompts.turnNudge, compressedBlockGuidance) for (const message of messages) { if (!state.nudges.turnNudgeAnchors.has(message.info.id)) continue @@ -259,8 +289,91 @@ export function applyAnchoredNudges( } } - applyAnchoredNudge(turnNudgeAnchors, messages, turnNudge) + return turnNudgeAnchors +} - const iterationNudge = appendGuidanceToDcpTag(prompts.iterationNudge, compressedBlockGuidance) - applyAnchoredNudge(state.nudges.iterationNudgeAnchors, messages, iterationNudge) +function applyRangeModeAnchoredNudge( + anchorMessageIds: Set, + messages: WithParts[], + baseNudgeText: string, + compressedBlockGuidance: string, +): void { + const nudgeText = appendGuidanceToDcpTag(baseNudgeText, compressedBlockGuidance) + if (!nudgeText.trim()) { + return + } + + for (const { message } of collectAnchoredMessages(anchorMessageIds, messages)) { + injectAnchoredNudge(message, nudgeText) + } +} + +function applyMessageModeAnchoredNudge( + anchorMessageIds: Set, + messages: WithParts[], + baseNudgeText: string, + compressionPriorities?: CompressionPriorityMap, +): void { + for (const { message, index } of collectAnchoredMessages(anchorMessageIds, messages)) { + const priorityGuidance = buildMessagePriorityGuidance( + messages, + compressionPriorities, + index, + MESSAGE_MODE_NUDGE_PRIORITY, + ) + const nudgeText = appendGuidanceToDcpTag(baseNudgeText, priorityGuidance) + injectAnchoredNudge(message, nudgeText) + } +} + +export function applyAnchoredNudges( + state: SessionState, + config: PluginConfig, + messages: WithParts[], + prompts: RuntimePrompts, + compressionPriorities?: CompressionPriorityMap, +): void { + const turnNudgeAnchors = collectTurnNudgeAnchors(state, config, messages) + + if (config.compress.mode === "message") { + applyMessageModeAnchoredNudge( + state.nudges.contextLimitAnchors, + messages, + prompts.contextLimitNudge, + compressionPriorities, + ) + applyMessageModeAnchoredNudge( + turnNudgeAnchors, + messages, + prompts.turnNudge, + compressionPriorities, + ) + applyMessageModeAnchoredNudge( + state.nudges.iterationNudgeAnchors, + messages, + prompts.iterationNudge, + compressionPriorities, + ) + return + } + + const compressedBlockGuidance = buildCompressedBlockGuidance(state) + applyRangeModeAnchoredNudge( + state.nudges.contextLimitAnchors, + messages, + prompts.contextLimitNudge, + compressedBlockGuidance, + ) + applyRangeModeAnchoredNudge( + turnNudgeAnchors, + messages, + prompts.turnNudge, + compressedBlockGuidance, + ) + applyRangeModeAnchoredNudge( + state.nudges.iterationNudgeAnchors, + messages, + prompts.iterationNudge, + compressedBlockGuidance, + ) } diff --git a/lib/messages/priority.ts b/lib/messages/priority.ts new file mode 100644 index 00000000..cc95ef16 --- /dev/null +++ b/lib/messages/priority.ts @@ -0,0 +1,102 @@ +import type { PluginConfig } from "../config" +import { countAllMessageTokens } from "../strategies/utils" +import { isMessageCompacted } from "../shared-utils" +import type { SessionState, WithParts } from "../state" +import { isIgnoredUserMessage, isProtectedUserMessage } from "./utils" + +const MEDIUM_PRIORITY_MIN_TOKENS = 500 +const HIGH_PRIORITY_MIN_TOKENS = 5000 + +export type MessagePriority = "low" | "medium" | "high" + +export interface CompressionPriorityEntry { + ref: string + tokenCount: number + priority: MessagePriority +} + +export type CompressionPriorityMap = Map + +export function buildPriorityMap( + config: PluginConfig, + state: SessionState, + messages: WithParts[], +): CompressionPriorityMap { + if (config.compress.mode !== "message") { + return new Map() + } + const priorities: CompressionPriorityMap = new Map() + + for (const message of messages) { + if (message.info.role === "user" && isIgnoredUserMessage(message)) { + continue + } + + if (isProtectedUserMessage(config, message)) { + continue + } + + if (isMessageCompacted(state, message)) { + continue + } + + const rawMessageId = message.info.id + if (typeof rawMessageId !== "string" || rawMessageId.length === 0) { + continue + } + + const ref = state.messageIds.byRawId.get(rawMessageId) + if (!ref) { + continue + } + + const tokenCount = countAllMessageTokens(message) + priorities.set(rawMessageId, { + ref, + tokenCount, + priority: classifyMessagePriority(tokenCount), + }) + } + + return priorities +} + +export function classifyMessagePriority(tokenCount: number): MessagePriority { + if (tokenCount >= HIGH_PRIORITY_MIN_TOKENS) { + return "high" + } + + if (tokenCount >= MEDIUM_PRIORITY_MIN_TOKENS) { + return "medium" + } + + return "low" +} + +export function listPriorityRefsBeforeIndex( + messages: WithParts[], + priorities: CompressionPriorityMap, + anchorIndex: number, + priority: MessagePriority, +): string[] { + const refs: string[] = [] + const seen = new Set() + const upperBound = Math.max(0, Math.min(anchorIndex, messages.length)) + + for (let index = 0; index < upperBound; index++) { + const rawMessageId = messages[index]?.info.id + if (typeof rawMessageId !== "string") { + continue + } + + const entry = priorities.get(rawMessageId) + if (!entry || entry.priority !== priority || seen.has(entry.ref)) { + continue + } + + seen.add(entry.ref) + refs.push(entry.ref) + } + + return refs +} diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts index 62353969..fa5b7098 100644 --- a/lib/messages/prune.ts +++ b/lib/messages/prune.ts @@ -2,7 +2,7 @@ import type { SessionState, WithParts } from "../state" import type { Logger } from "../logger" import type { PluginConfig } from "../config" import { isMessageCompacted, getLastUserMessage } from "../shared-utils" -import { createSyntheticUserMessage } from "./utils" +import { createSyntheticUserMessage, replaceBlockIdsWithBlocked } from "./utils" import type { UserMessage } from "@opencode-ai/sdk/v2" const PRUNED_TOOL_OUTPUT_REPLACEMENT = @@ -16,7 +16,7 @@ export const prune = ( config: PluginConfig, messages: WithParts[], ): void => { - filterCompressedRanges(state, logger, messages) + filterCompressedRanges(state, logger, config, messages) // pruneFullTool(state, logger, messages) pruneToolOutputs(state, logger, messages) pruneToolInputs(state, logger, messages) @@ -158,6 +158,7 @@ const pruneToolErrors = (state: SessionState, logger: Logger, messages: WithPart const filterCompressedRanges = ( state: SessionState, logger: Logger, + config: PluginConfig, messages: WithParts[], ): void => { if ( @@ -194,7 +195,10 @@ const filterCompressedRanges = ( if (userMessage) { const userInfo = userMessage.info as UserMessage - const summaryContent = rawSummaryContent + const summaryContent = + config.compress.mode === "message" + ? replaceBlockIdsWithBlocked(rawSummaryContent) + : rawSummaryContent const summarySeed = `${summary.blockId}:${summary.anchorMessageId}` result.push( createSyntheticUserMessage( diff --git a/lib/messages/utils.ts b/lib/messages/utils.ts index 03ce5690..8ccf6561 100644 --- a/lib/messages/utils.ts +++ b/lib/messages/utils.ts @@ -1,11 +1,15 @@ import { createHash } from "node:crypto" +import type { PluginConfig } from "../config" import { isMessageCompacted } from "../shared-utils" import type { SessionState, WithParts } from "../state" import type { UserMessage } from "@opencode-ai/sdk/v2" const SUMMARY_ID_HASH_LENGTH = 16 -const DCP_MESSAGE_ID_TAG_REGEX = /(?:m\d+|b\d+)<\/dcp-message-id>/g -const DCP_SYSTEM_REMINDER_REGEX = /]*>[\s\S]*?<\/dcp-system-reminder>/g +const DCP_MESSAGE_ID_TAG_REGEX = + /])[^>]*>(?:m\d+|b\d+|BLOCKED)<\/dcp-message-id>/g +const DCP_BLOCK_ID_TAG_REGEX = /(])[^>]*>)b\d+(<\/dcp-message-id>)/g +const DCP_SYSTEM_REMINDER_REGEX = + /])[^>]*>[\s\S]*?<\/dcp-system-reminder>/g const generateStableId = (prefix: string, seed: string): string => { const hash = createHash("sha256").update(seed).digest("hex").slice(0, SUMMARY_ID_HASH_LENGTH) @@ -66,6 +70,35 @@ export const createSyntheticTextPart = ( type MessagePart = WithParts["parts"][number] type ToolPart = Extract +type TextPart = Extract + +const findLastTextPart = (message: WithParts): TextPart | null => { + for (let i = message.parts.length - 1; i >= 0; i--) { + const part = message.parts[i] + if (part.type === "text") { + return part + } + } + + return null +} + +export const appendToTextPart = (message: WithParts, injection: string): boolean => { + const textPart = findLastTextPart(message) + if (!textPart || typeof textPart.text !== "string") { + return false + } + + const normalizedInjection = injection.replace(/^\n+/, "") + if (!normalizedInjection.trim()) { + return false + } + + const baseText = textPart.text.replace(/\n*$/, "") + textPart.text = + baseText.length > 0 ? `${baseText}\n\n${normalizedInjection}` : normalizedInjection + return true +} export const appendIdToTool = (part: ToolPart, tag: string): boolean => { if (part.type !== "tool") { @@ -127,6 +160,19 @@ export const isIgnoredUserMessage = (message: WithParts): boolean => { return true } +export function isProtectedUserMessage(config: PluginConfig, message: WithParts): boolean { + return ( + config.compress.mode === "message" && + config.compress.protectUserMessages && + message.info.role === "user" && + !isIgnoredUserMessage(message) + ) +} + +export const replaceBlockIdsWithBlocked = (text: string): string => { + return text.replace(DCP_BLOCK_ID_TAG_REGEX, "$1BLOCKED$2") +} + export const stripHallucinationsFromString = (text: string): string => { return text.replace(DCP_SYSTEM_REMINDER_REGEX, "").replace(DCP_MESSAGE_ID_TAG_REGEX, "") } diff --git a/lib/prompts/compress-message.ts b/lib/prompts/compress-message.ts new file mode 100644 index 00000000..9a5f4911 --- /dev/null +++ b/lib/prompts/compress-message.ts @@ -0,0 +1,43 @@ +export const COMPRESS_MESSAGE = `Collapse selected individual messages in the conversation into detailed summaries. + +THE SUMMARY +Your summary must be EXHAUSTIVE. Capture file paths, function signatures, decisions made, constraints discovered, key findings, tool outcomes, and user intent details that matter... EVERYTHING that preserves the value of the selected message after the raw message is removed. + +USER INTENT FIDELITY +When a selected message contains user intent, preserve that intent with extra care. Do not change scope, constraints, priorities, acceptance criteria, or requested outcomes. +Directly quote short user instructions when that best preserves exact meaning. + +Yet be LEAN. Strip away the noise: failed attempts that led nowhere, verbose tool output, and repetition. What remains should be pure signal - golden nuggets of detail that preserve full understanding with zero ambiguity. + +MESSAGE IDS +You specify individual raw messages by ID using the injected IDs visible in the conversation: + +- \`mNNNN\` IDs identify raw messages + +Each message has an ID inside XML metadata tags like \`m0007\`. +The ID tag appears at the end of the message it belongs to — it identifies the message above it, not the one below it. +Treat these tags as message metadata only, not as content to summarize. Use only the inner \`mNNNN\` value as the \`messageId\`. +The \`priority\` attribute indicates relative context cost. Prefer higher-priority closed messages before lower-priority ones. +Messages marked as \`BLOCKED\` cannot be compressed. + +Rules: + +- Pick each \`messageId\` directly from injected IDs visible in context. +- Only use raw message IDs of the form \`mNNNN\`. +- Ignore XML attributes such as \`priority\` when copying the ID; use only the inner \`mNNNN\` value. +- Do NOT use compressed block IDs like \`bN\`. +- Do not invent IDs. Use only IDs that are present in context. +- Do not target prior compressed blocks or block summaries. + +BATCHING +Select MANY messages in a single tool call when they are independently safe to compress. +Each entry should summarize exactly one message, and the tool can receive as many entries as needed in one batch. +When several messages are equally safe to compress, prefer higher-priority messages first. + +Because each message is compressed independently: + +- Do not describe ranges +- Do not use start/end boundaries +- Do not use compressed block placeholders +- Do not reference prior compressed blocks with \`(bN)\` +` diff --git a/lib/prompts/compress.ts b/lib/prompts/compress-range.ts similarity index 63% rename from lib/prompts/compress.ts rename to lib/prompts/compress-range.ts index 8bc2e34f..bfc312f8 100644 --- a/lib/prompts/compress.ts +++ b/lib/prompts/compress-range.ts @@ -1,9 +1,4 @@ -export const COMPRESS = `Collapse a range in the conversation into a detailed summary. - -THE PHILOSOPHY OF COMPRESS -\`compress\` transforms verbose conversation sequences into dense, high-fidelity summaries. This is not cleanup - it is crystallization. Your summary becomes the authoritative record of what transpired. - -Think of compression as phase transitions: raw exploration becomes refined understanding. The original context served its purpose; your summary now carries that understanding forward. +export const COMPRESS_RANGE = `Collapse a range in the conversation into a detailed summary. THE SUMMARY Your summary must be EXHAUSTIVE. Capture file paths, function signatures, decisions made, constraints discovered, key findings... EVERYTHING that maintains context integrity. This is not a brief note - it is an authoritative record so faithful that the original conversation adds no value. @@ -43,25 +38,6 @@ When you use compressed block placeholders, write the surrounding summary text s - Do not write text that depends on the placeholder staying literal (for example, "as noted in \`(b2)\`"). - Your final meaning must be coherent once each placeholder is replaced with its full compressed block content. -THE WAYS OF COMPRESS -Compress when a range is genuinely closed and the raw conversation has served its purpose: - -Research concluded and findings are clear -Implementation finished and verified -Exploration exhausted and patterns understood - -Compress smaller ranges when: -You need to discard dead-end noise without waiting for a whole chapter to close -You need to preserve key findings from a narrow slice while freeing context quickly -You can bound a stale range cleanly with injected IDs - -Do NOT compress when: -You may need exact code, error messages, or file contents from the range in the immediate next steps -Work in that area is still active or likely to resume immediately -You cannot identify reliable boundaries yet - -Before compressing, ask: _"Is this range closed enough to become summary-only right now?"_ Compression is irreversible. The summary replaces everything in the range. - BOUNDARY IDS You specify boundaries by ID using the injected IDs visible in the conversation: @@ -69,6 +45,7 @@ You specify boundaries by ID using the injected IDs visible in the conversation: - \`bN\` IDs identify previously compressed blocks Each message has an ID inside XML metadata tags like \`...\`. +The ID tag appears at the end of the message it belongs to — it identifies the message above it, not the one below it. Treat these tags as boundary metadata only, not as tool result content. Rules: @@ -76,9 +53,8 @@ Rules: - Pick \`startId\` and \`endId\` directly from injected IDs in context. - IDs must exist in the current visible context. - \`startId\` must appear before \`endId\`. -- Prefer boundaries that produce short, closed ranges. - Do not invent IDs. Use only IDs that are present in context. -PARALLEL COMPRESS EXECUTION -When multiple independent ranges are ready and their boundaries do not overlap, launch MULTIPLE \`compress\` calls in parallel in a single response. This is the PREFERRED pattern over a single large-range compression when the work can be safely split. Run compression sequentially only when ranges overlap or when a later range depends on the result of an earlier compression. +BATCHING +When multiple independent ranges are ready and their boundaries do not overlap, include all of them as separate entries in the \`content\` array of a single tool call. Each entry should have its own \`startId\`, \`endId\`, and \`summary\`. ` diff --git a/lib/prompts/internal-overlays.ts b/lib/prompts/internal-overlays.ts index e445dafd..e0dd317f 100644 --- a/lib/prompts/internal-overlays.ts +++ b/lib/prompts/internal-overlays.ts @@ -18,28 +18,34 @@ All subsequent messages in the session will have IDs. ` -export const NESTED_FORMAT_OVERLAY = ` +export const RANGE_FORMAT_OVERLAY = ` THE FORMAT OF COMPRESS \`\`\` { topic: string, // Short label (3-5 words) - e.g., "Auth System Exploration" - content: { - startId: string, // Boundary ID at range start: mNNNN or bN - endId: string, // Boundary ID at range end: mNNNN or bN - summary: string // Complete technical summary replacing all content in range - } + content: [ // One or more ranges to compress + { + startId: string, // Boundary ID at range start: mNNNN or bN + endId: string, // Boundary ID at range end: mNNNN or bN + summary: string // Complete technical summary replacing all content in range + } + ] } \`\`\`` -export const FLAT_FORMAT_OVERLAY = ` +export const MESSAGE_FORMAT_OVERLAY = ` THE FORMAT OF COMPRESS \`\`\` { - topic: string, // Short label (3-5 words) - e.g., "Auth System Exploration" - startId: string, // Boundary ID at range start: mNNNN or bN - endId: string, // Boundary ID at range end: mNNNN or bN - summary: string // Complete technical summary replacing all content in range + topic: string, // Short label (3-5 words) for the overall batch + content: [ // One or more messages to compress independently + { + messageId: string, // Raw message ID only: mNNNN (ignore metadata attributes like priority) + topic: string, // Short label (3-5 words) for this one message summary + summary: string // Complete technical summary replacing that one message + } + ] } \`\`\`` diff --git a/lib/prompts/message-priority-guidance.ts b/lib/prompts/message-priority-guidance.ts new file mode 100644 index 00000000..f41ab626 --- /dev/null +++ b/lib/prompts/message-priority-guidance.ts @@ -0,0 +1,9 @@ +export function renderMessagePriorityGuidance(priorityLabel: string, refs: string[]): string { + const refList = refs.length > 0 ? refs.join(", ") : "none" + + return [ + "Message priority context:", + "- Higher-priority older messages consume more context and should be compressed before lower-priority ones when safely closed.", + `- ${priorityLabel}-priority message IDs before this point: ${refList}`, + ].join("\n") +} diff --git a/lib/prompts/store.ts b/lib/prompts/store.ts index d088f07b..885ae670 100644 --- a/lib/prompts/store.ts +++ b/lib/prompts/store.ts @@ -3,7 +3,8 @@ import { join, dirname } from "path" import { homedir } from "os" import type { Logger } from "../logger" import { SYSTEM as SYSTEM_PROMPT } from "./system" -import { COMPRESS as COMPRESS_PROMPT } from "./compress" +import { COMPRESS_RANGE as COMPRESS_RANGE_PROMPT } from "./compress-range" +import { COMPRESS_MESSAGE as COMPRESS_MESSAGE_PROMPT } from "./compress-message" import { CONTEXT_LIMIT_NUDGE } from "./context-limit-nudge" import { TURN_NUDGE } from "./turn-nudge" import { ITERATION_NUDGE } from "./iteration-nudge" @@ -11,14 +12,16 @@ import { MANUAL_MODE_SYSTEM_OVERLAY, SUBAGENT_SYSTEM_OVERLAY } from "./internal- export type PromptKey = | "system" - | "compress" + | "compress-range" + | "compress-message" | "context-limit-nudge" | "turn-nudge" | "iteration-nudge" type EditablePromptField = | "system" - | "compress" + | "compressRange" + | "compressMessage" | "contextLimitNudge" | "turnNudge" | "iterationNudge" @@ -45,7 +48,8 @@ interface PromptPaths { export interface RuntimePrompts { system: string - compress: string + compressRange: string + compressMessage: string contextLimitNudge: string turnNudge: string iterationNudge: string @@ -63,12 +67,20 @@ const PROMPT_DEFINITIONS: PromptDefinition[] = [ runtimeField: "system", }, { - key: "compress", - fileName: "compress.md", - label: "Compress", - description: "compress tool instructions and summary constraints", - usage: "Registered as the compress tool description", - runtimeField: "compress", + key: "compress-range", + fileName: "compress-range.md", + label: "Compress Range", + description: "range-mode compress tool instructions and summary constraints", + usage: "Registered as the range-mode compress tool description", + runtimeField: "compressRange", + }, + { + key: "compress-message", + fileName: "compress-message.md", + label: "Compress Message", + description: "message-mode compress tool instructions and summary constraints", + usage: "Registered as the message-mode compress tool description", + runtimeField: "compressMessage", }, { key: "context-limit-nudge", @@ -98,7 +110,8 @@ const PROMPT_DEFINITIONS: PromptDefinition[] = [ export const PROMPT_KEYS: PromptKey[] = [ "system", - "compress", + "compress-range", + "compress-message", "context-limit-nudge", "turn-nudge", "iteration-nudge", @@ -112,7 +125,8 @@ const DEFAULTS_README_FILE = "README.md" const BUNDLED_EDITABLE_PROMPTS: Record = { system: SYSTEM_PROMPT, - compress: COMPRESS_PROMPT, + compressRange: COMPRESS_RANGE_PROMPT, + compressMessage: COMPRESS_MESSAGE_PROMPT, contextLimitNudge: CONTEXT_LIMIT_NUDGE, turnNudge: TURN_NUDGE, iterationNudge: ITERATION_NUDGE, @@ -126,7 +140,8 @@ const INTERNAL_PROMPT_OVERLAYS = { function createBundledRuntimePrompts(): RuntimePrompts { return { system: BUNDLED_EDITABLE_PROMPTS.system, - compress: BUNDLED_EDITABLE_PROMPTS.compress, + compressRange: BUNDLED_EDITABLE_PROMPTS.compressRange, + compressMessage: BUNDLED_EDITABLE_PROMPTS.compressMessage, contextLimitNudge: BUNDLED_EDITABLE_PROMPTS.contextLimitNudge, turnNudge: BUNDLED_EDITABLE_PROMPTS.turnNudge, iterationNudge: BUNDLED_EDITABLE_PROMPTS.iterationNudge, @@ -232,7 +247,7 @@ function toEditablePromptText(definition: PromptDefinition, rawContent: string): normalized = stripConditionalTag(normalized, "subagent") } - if (definition.key !== "compress") { + if (definition.key !== "compress-range" && definition.key !== "compress-message") { normalized = normalizeReminderPromptContent(normalized) } @@ -245,7 +260,7 @@ function wrapRuntimePromptContent(definition: PromptDefinition, editableText: st return "" } - if (definition.key === "compress") { + if (definition.key === "compress-range" || definition.key === "compress-message") { return trimmed } diff --git a/lib/prompts/system.ts b/lib/prompts/system.ts index 984f1e5a..69ffbb80 100644 --- a/lib/prompts/system.ts +++ b/lib/prompts/system.ts @@ -1,29 +1,46 @@ export const SYSTEM = ` You operate in a context-constrained environment. Manage context continuously to avoid buildup and preserve retrieval quality. Efficient context management is paramount for your agentic performance. -The ONLY tool you have for context management is \`compress\`. It replaces a contiguous portion of the conversation (inclusive) with a technical summary you produce. +The ONLY tool you have for context management is \`compress\`. It replaces older conversation content with technical summaries you produce. \`\` and \`\` tags are environment-injected metadata. Do not output them. +THE PHILOSOPHY OF COMPRESS +\`compress\` transforms conversation content into dense, high-fidelity summaries. This is not cleanup - it is crystallization. Your summary becomes the authoritative record of what transpired. + +Think of compression as phase transitions: raw exploration becomes refined understanding. The original context served its purpose; your summary now carries that understanding forward. + OPERATING STANCE -Prefer short, closed, summary-safe ranges. -When multiple independent stale ranges exist, prefer several short compressions (in parallel when possible) over one large-range compression. +Prefer short, closed, summary-safe compressions. +When multiple independent stale sections exist, prefer several focused compressions (in parallel when possible) over one broad compression. Use \`compress\` as steady housekeeping while you work. CADENCE, SIGNALS, AND LATENCY - No fixed threshold mandates compression -- Prioritize closedness and independence over raw range size +- Prioritize closedness and independence over raw size - Prefer smaller, regular compressions over infrequent massive compressions for better latency and summary quality -- When multiple independent stale ranges are ready, batch compressions in parallel +- When multiple independent stale sections are ready, batch compressions in parallel + +COMPRESS WHEN + +A section is genuinely closed and the raw conversation has served its purpose: + +- Research concluded and findings are clear +- Implementation finished and verified +- Exploration exhausted and patterns understood +- Dead-end noise can be discarded without waiting for a whole chapter to close DO NOT COMPRESS IF -- raw context is still relevant and needed for edits or precise references -- the task in the target range is still actively in progress +- Raw context is still relevant and needed for edits or precise references +- The target content is still actively in progress +- You may need exact code, error messages, or file contents in the immediate next steps + +Before compressing, ask: _"Is this section closed enough to become summary-only right now?"_ -Evaluate conversation signal-to-noise REGULARLY. Use \`compress\` deliberately with quality-first summaries. Prefer multiple short, independent range compressions before considering broader ranges, and prioritize ranges intelligently to maintain a high-signal context window that supports your agency +Evaluate conversation signal-to-noise REGULARLY. Use \`compress\` deliberately with quality-first summaries. Prioritize stale content intelligently to maintain a high-signal context window that supports your agency. -It is of your responsibility to keep a sharp, high-quality context window for optimal performance +It is of your responsibility to keep a sharp, high-quality context window for optimal performance. ` diff --git a/lib/state/persistence.ts b/lib/state/persistence.ts index 8a62c432..2431eee3 100644 --- a/lib/state/persistence.ts +++ b/lib/state/persistence.ts @@ -18,6 +18,7 @@ export interface PersistedPruneMessagesState { activeBlockIds: number[] activeByAnchorMessageId: Record nextBlockId: number + nextRunId: number } export interface PersistedPrune { @@ -85,6 +86,7 @@ export async function saveSessionState( sessionState.prune.messages.activeByAnchorMessageId, ), nextBlockId: sessionState.prune.messages.nextBlockId, + nextRunId: sessionState.prune.messages.nextRunId, }, }, nudges: { diff --git a/lib/state/types.ts b/lib/state/types.ts index 7424c9e9..6fe2fb9e 100644 --- a/lib/state/types.ts +++ b/lib/state/types.ts @@ -27,12 +27,17 @@ export interface PrunedMessageEntry { activeBlockIds: number[] } +export type CompressionMode = "range" | "message" + export interface CompressionBlock { blockId: number + runId: number active: boolean deactivatedByUser: boolean compressedTokens: number + mode?: CompressionMode topic: string + batchTopic?: string startId: string endId: string anchorMessageId: string @@ -56,6 +61,7 @@ export interface PruneMessagesState { activeBlockIds: Set activeByAnchorMessageId: Map nextBlockId: number + nextRunId: number } export interface Prune { diff --git a/lib/state/utils.ts b/lib/state/utils.ts index cd1f0662..4e018278 100644 --- a/lib/state/utils.ts +++ b/lib/state/utils.ts @@ -14,6 +14,7 @@ interface PersistedPruneMessagesState { activeBlockIds?: number[] activeByAnchorMessageId?: Record nextBlockId?: number + nextRunId?: number } export async function isSubAgentSession(client: any, sessionID: string): Promise { @@ -70,6 +71,7 @@ export function createPruneMessagesState(): PruneMessagesState { activeBlockIds: new Set(), activeByAnchorMessageId: new Map(), nextBlockId: 1, + nextRunId: 1, } } @@ -84,6 +86,9 @@ export function loadPruneMessagesState( if (typeof persisted.nextBlockId === "number" && Number.isInteger(persisted.nextBlockId)) { state.nextBlockId = Math.max(1, persisted.nextBlockId) } + if (typeof persisted.nextRunId === "number" && Number.isInteger(persisted.nextRunId)) { + state.nextRunId = Math.max(1, persisted.nextRunId) + } if (persisted.byMessageId && typeof persisted.byMessageId === "object") { for (const [messageId, entry] of Object.entries(persisted.byMessageId)) { @@ -143,6 +148,12 @@ export function loadPruneMessagesState( state.blocksById.set(blockId, { blockId, + runId: + typeof block.runId === "number" && + Number.isInteger(block.runId) && + block.runId > 0 + ? block.runId + : blockId, active: block.active === true, deactivatedByUser: block.deactivatedByUser === true, compressedTokens: @@ -150,7 +161,14 @@ export function loadPruneMessagesState( Number.isFinite(block.compressedTokens) ? Math.max(0, block.compressedTokens) : 0, + mode: block.mode === "range" || block.mode === "message" ? block.mode : undefined, topic: typeof block.topic === "string" ? block.topic : "", + batchTopic: + typeof block.batchTopic === "string" + ? block.batchTopic + : typeof block.topic === "string" + ? block.topic + : "", startId: typeof block.startId === "string" ? block.startId : "", endId: typeof block.endId === "string" ? block.endId : "", anchorMessageId: @@ -210,6 +228,9 @@ export function loadPruneMessagesState( if (blockId >= state.nextBlockId) { state.nextBlockId = blockId + 1 } + if (block.runId >= state.nextRunId) { + state.nextRunId = block.runId + 1 + } } return state diff --git a/lib/strategies/index.ts b/lib/strategies/index.ts index f8922df9..d6ef1009 100644 --- a/lib/strategies/index.ts +++ b/lib/strategies/index.ts @@ -1,3 +1,2 @@ export { deduplicate } from "./deduplication" -export { supersedeWrites } from "./supersede-writes" export { purgeErrors } from "./purge-errors" diff --git a/lib/strategies/supersede-writes.ts b/lib/strategies/supersede-writes.ts deleted file mode 100644 index 0f07ffca..00000000 --- a/lib/strategies/supersede-writes.ts +++ /dev/null @@ -1,115 +0,0 @@ -import { PluginConfig } from "../config" -import { Logger } from "../logger" -import type { SessionState, WithParts } from "../state" -import { getFilePathsFromParameters, isFilePathProtected } from "../protected-patterns" -import { getTotalToolTokens } from "./utils" - -/** - * Supersede Writes strategy - prunes write tool inputs for files that have - * subsequently been read. When a file is written and later read, the original - * write content becomes redundant since the current file state is captured - * in the read result. - * - * Modifies the session state in place to add pruned tool call IDs. - */ -export const supersedeWrites = ( - state: SessionState, - logger: Logger, - config: PluginConfig, - messages: WithParts[], -): void => { - if (state.manualMode && !config.manualMode.automaticStrategies) { - return - } - - if (!config.strategies.supersedeWrites.enabled) { - return - } - - const allToolIds = state.toolIdList - if (allToolIds.length === 0) { - return - } - - // Filter out IDs already pruned - const unprunedIds = allToolIds.filter((id) => !state.prune.tools.has(id)) - if (unprunedIds.length === 0) { - return - } - - // Track write tools by file path: filePath -> [{ id, index }] - // We track index to determine chronological order - const writesByFile = new Map() - - // Track read file paths with their index - const readsByFile = new Map() - - for (let i = 0; i < allToolIds.length; i++) { - const id = allToolIds[i] - const metadata = state.toolParameters.get(id) - if (!metadata) { - continue - } - - const filePaths = getFilePathsFromParameters(metadata.tool, metadata.parameters) - if (filePaths.length === 0) { - continue - } - const filePath = filePaths[0] - - if (isFilePathProtected(filePaths, config.protectedFilePatterns)) { - continue - } - - if (metadata.tool === "write") { - if (!writesByFile.has(filePath)) { - writesByFile.set(filePath, []) - } - const writes = writesByFile.get(filePath) - if (writes) { - writes.push({ id, index: i }) - } - } else if (metadata.tool === "read") { - if (!readsByFile.has(filePath)) { - readsByFile.set(filePath, []) - } - const reads = readsByFile.get(filePath) - if (reads) { - reads.push(i) - } - } - } - - // Find writes that are superseded by subsequent reads - const newPruneIds: string[] = [] - - for (const [filePath, writes] of writesByFile.entries()) { - const reads = readsByFile.get(filePath) - if (!reads || reads.length === 0) { - continue - } - - // For each write, check if there's a read that comes after it - for (const write of writes) { - // Skip if already pruned - if (state.prune.tools.has(write.id)) { - continue - } - - // Check if any read comes after this write - const hasSubsequentRead = reads.some((readIndex) => readIndex > write.index) - if (hasSubsequentRead) { - newPruneIds.push(write.id) - } - } - } - - if (newPruneIds.length > 0) { - state.stats.totalPruneTokens += getTotalToolTokens(state, newPruneIds) - for (const id of newPruneIds) { - const entry = state.toolParameters.get(id) - state.prune.tools.set(id, entry?.tokenCount ?? 0) - } - logger.debug(`Marked ${newPruneIds.length} superseded write tool calls for pruning`) - } -} diff --git a/lib/tools/compress.ts b/lib/tools/compress.ts deleted file mode 100644 index 6d2b62f1..00000000 --- a/lib/tools/compress.ts +++ /dev/null @@ -1,224 +0,0 @@ -import { tool } from "@opencode-ai/plugin" -import type { ToolContext } from "./types" -import { ensureSessionInitialized } from "../state" -import { - appendMissingBlockSummaries, - appendProtectedUserMessages, - appendProtectedTools, - wrapCompressedSummary, - allocateBlockId, - applyCompressionState, - buildSearchContext, - fetchSessionMessages, - COMPRESSED_BLOCK_HEADER, - injectBlockPlaceholders, - parseBlockPlaceholders, - resolveAnchorMessageId, - resolveBoundaryIds, - resolveRange, - normalizeCompressArgs, - validateCompressArgs, - validateSummaryPlaceholders, - type CompressToolArgs, -} from "./utils" -import { isIgnoredUserMessage } from "../messages/utils" -import { assignMessageRefs } from "../message-ids" -import { getCurrentParams, getCurrentTokenUsage, countTokens } from "../strategies/utils" -import { deduplicate, supersedeWrites, purgeErrors } from "../strategies" -import { saveSessionState } from "../state/persistence" -import { sendCompressNotification } from "../ui/notification" -import { NESTED_FORMAT_OVERLAY, FLAT_FORMAT_OVERLAY } from "../prompts/internal-overlays" - -// This schema looks better in the TUI (non primitive args aren't displayed), but LLMs are more likely to fail -// the tool call -function buildNestedSchema() { - return { - topic: tool.schema - .string() - .describe("Short label (3-5 words) for display - e.g., 'Auth System Exploration'"), - content: tool.schema - .object({ - startId: tool.schema - .string() - .describe( - "Message or block ID marking the beginning of range (e.g. m0001, b2)", - ), - endId: tool.schema - .string() - .describe("Message or block ID marking the end of range (e.g. m0012, b5)"), - summary: tool.schema - .string() - .describe("Complete technical summary replacing all content in range"), - }) - .describe("Compression details: ID boundaries and replacement summary"), - } -} - -// Simpler schema for models that are not as good at tool calling reliably -function buildFlatSchema() { - return { - topic: tool.schema - .string() - .describe("Short label (3-5 words) for display - e.g., 'Auth System Exploration'"), - startId: tool.schema - .string() - .describe("Message or block ID marking the beginning of range (e.g. m0001, b2)"), - endId: tool.schema - .string() - .describe("Message or block ID marking the end of range (e.g. m0012, b5)"), - summary: tool.schema - .string() - .describe("Complete technical summary replacing all content in range"), - } -} - -export function createCompressTool(ctx: ToolContext): ReturnType { - ctx.prompts.reload() - const runtimePrompts = ctx.prompts.getRuntimePrompts() - const useFlatSchema = ctx.config.compress.flatSchema - - return tool({ - description: - runtimePrompts.compress + (useFlatSchema ? FLAT_FORMAT_OVERLAY : NESTED_FORMAT_OVERLAY), - args: useFlatSchema ? buildFlatSchema() : buildNestedSchema(), - async execute(args, toolCtx) { - if (ctx.state.manualMode && ctx.state.manualMode !== "compress-pending") { - throw new Error( - "Manual mode: compress blocked. Do not retry until `` appears in user context.", - ) - } - - await toolCtx.ask({ - permission: "compress", - patterns: ["*"], - always: ["*"], - metadata: {}, - }) - - const compressArgs = normalizeCompressArgs(args as Record) - validateCompressArgs(compressArgs) - - toolCtx.metadata({ - title: `Compress: ${compressArgs.topic}`, - }) - - const rawMessages = await fetchSessionMessages(ctx.client, toolCtx.sessionID) - - await ensureSessionInitialized( - ctx.client, - ctx.state, - toolCtx.sessionID, - ctx.logger, - rawMessages, - ctx.config.manualMode.enabled, - ) - - assignMessageRefs(ctx.state, rawMessages) - - deduplicate(ctx.state, ctx.logger, ctx.config, rawMessages) - // supersedeWrites(ctx.state, ctx.logger, ctx.config, rawMessages) - purgeErrors(ctx.state, ctx.logger, ctx.config, rawMessages) - - const searchContext = buildSearchContext(ctx.state, rawMessages) - - const { startReference, endReference } = resolveBoundaryIds( - searchContext, - ctx.state, - compressArgs.content.startId, - compressArgs.content.endId, - ) - - const range = resolveRange(searchContext, startReference, endReference) - const anchorMessageId = resolveAnchorMessageId(range.startReference) - - const parsedPlaceholders = parseBlockPlaceholders(compressArgs.content.summary) - const missingRequiredBlockIds = validateSummaryPlaceholders( - parsedPlaceholders, - range.requiredBlockIds, - range.startReference, - range.endReference, - searchContext.summaryByBlockId, - ) - - const injected = injectBlockPlaceholders( - compressArgs.content.summary, - parsedPlaceholders, - searchContext.summaryByBlockId, - range.startReference, - range.endReference, - ) - - const summaryWithUserMessages = appendProtectedUserMessages( - injected.expandedSummary, - range, - searchContext, - ctx.state, - ctx.config.compress.protectUserMessages, - ) - - const summaryWithProtectedTools = await appendProtectedTools( - ctx.client, - ctx.state, - ctx.config.experimental.allowSubAgents, - summaryWithUserMessages, - range, - searchContext, - ctx.config.compress.protectedTools, - ctx.config.protectedFilePatterns, - ) - - const finalSummaryResult = appendMissingBlockSummaries( - summaryWithProtectedTools, - missingRequiredBlockIds, - searchContext.summaryByBlockId, - injected.consumedBlockIds, - ) - - const finalSummary = finalSummaryResult.expandedSummary - - const blockId = allocateBlockId(ctx.state) - const storedSummary = wrapCompressedSummary(blockId, finalSummary) - const summaryTokens = countTokens(storedSummary) - - const applied = applyCompressionState( - ctx.state, - { - topic: compressArgs.topic, - startId: compressArgs.content.startId, - endId: compressArgs.content.endId, - compressMessageId: toolCtx.messageID, - }, - range, - anchorMessageId, - blockId, - storedSummary, - finalSummaryResult.consumedBlockIds, - ) - - ctx.state.manualMode = ctx.state.manualMode ? "active" : false - await saveSessionState(ctx.state, ctx.logger) - - const params = getCurrentParams(ctx.state, rawMessages, ctx.logger) - const totalSessionTokens = getCurrentTokenUsage(rawMessages) - const sessionMessageIds = rawMessages - .filter((msg) => !(msg.info.role === "user" && isIgnoredUserMessage(msg))) - .map((msg) => msg.info.id) - - await sendCompressNotification( - ctx.client, - ctx.logger, - ctx.config, - ctx.state, - toolCtx.sessionID, - blockId, - compressArgs.content.summary, - summaryTokens, - totalSessionTokens, - sessionMessageIds, - params, - ) - - return `Compressed ${applied.messageIds.length} messages into ${COMPRESSED_BLOCK_HEADER}.` - }, - }) -} diff --git a/lib/tools/index.ts b/lib/tools/index.ts deleted file mode 100644 index 77f3cfcf..00000000 --- a/lib/tools/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export { ToolContext } from "./types" -export { createCompressTool } from "./compress" diff --git a/lib/tools/types.ts b/lib/tools/types.ts deleted file mode 100644 index 463f810d..00000000 --- a/lib/tools/types.ts +++ /dev/null @@ -1,13 +0,0 @@ -import type { SessionState } from "../state" -import type { PluginConfig } from "../config" -import type { Logger } from "../logger" -import type { PromptStore } from "../prompts/store" - -export interface ToolContext { - client: any - state: SessionState - logger: Logger - config: PluginConfig - workingDirectory: string - prompts: PromptStore -} diff --git a/lib/tools/utils.ts b/lib/tools/utils.ts deleted file mode 100644 index a3c24013..00000000 --- a/lib/tools/utils.ts +++ /dev/null @@ -1,975 +0,0 @@ -import type { CompressionBlock, SessionState, WithParts } from "../state" -import { formatBlockRef, formatMessageIdTag, parseBoundaryId } from "../message-ids" -import { isIgnoredUserMessage } from "../messages/utils" -import { countAllMessageTokens } from "../strategies/utils" -import { - getFilePathsFromParameters, - isFilePathProtected, - isToolNameProtected, -} from "../protected-patterns" -import { - buildSubagentResultText, - getSubAgentId, - mergeSubagentResult, -} from "../subagents/subagent-results" - -const BLOCK_PLACEHOLDER_REGEX = /\(b(\d+)\)|\{block_(\d+)\}/gi - -export interface CompressToolArgs { - topic: string - content: { - startId: string - endId: string - summary: string - } -} - -export interface FlatCompressToolArgs { - topic: string - startId: string - endId: string - summary: string -} - -export function normalizeCompressArgs(args: Record): CompressToolArgs { - if ("content" in args && typeof args.content === "object" && args.content !== null) { - return args as unknown as CompressToolArgs - } - - return { - topic: args.topic as string, - content: { - startId: args.startId as string, - endId: args.endId as string, - summary: args.summary as string, - }, - } -} - -export interface BoundaryReference { - kind: "message" | "compressed-block" - rawIndex: number - messageId?: string - blockId?: number - anchorMessageId?: string -} - -export interface SearchContext { - rawMessages: WithParts[] - rawMessagesById: Map - rawIndexById: Map - summaryByBlockId: Map -} - -export interface RangeResolution { - startReference: BoundaryReference - endReference: BoundaryReference - messageIds: string[] - messageTokenById: Map - toolIds: string[] - requiredBlockIds: number[] -} - -export interface ParsedBlockPlaceholder { - raw: string - blockId: number - startIndex: number - endIndex: number -} - -export interface InjectedSummaryResult { - expandedSummary: string - consumedBlockIds: number[] -} - -export interface AppliedCompressionResult { - compressedTokens: number - messageIds: string[] - newlyCompressedMessageIds: string[] - newlyCompressedToolIds: string[] -} - -export interface CompressionStateInput { - topic: string - startId: string - endId: string - compressMessageId: string -} - -export const COMPRESSED_BLOCK_HEADER = "[Compressed conversation section]" - -export function formatBlockPlaceholder(blockId: number): string { - return `(b${blockId})` -} - -export function validateCompressArgs(args: CompressToolArgs): void { - if (typeof args.topic !== "string" || args.topic.trim().length === 0) { - throw new Error("topic is required and must be a non-empty string") - } - - if (typeof args.content?.startId !== "string" || args.content.startId.trim().length === 0) { - throw new Error("content.startId is required and must be a non-empty string") - } - - if (typeof args.content?.endId !== "string" || args.content.endId.trim().length === 0) { - throw new Error("content.endId is required and must be a non-empty string") - } - - if (typeof args.content?.summary !== "string" || args.content.summary.trim().length === 0) { - throw new Error("content.summary is required and must be a non-empty string") - } -} - -export async function fetchSessionMessages(client: any, sessionId: string): Promise { - const response = await client.session.messages({ - path: { id: sessionId }, - }) - - const payload = (response?.data || response) as WithParts[] - return Array.isArray(payload) ? payload : [] -} - -export function buildSearchContext(state: SessionState, rawMessages: WithParts[]): SearchContext { - const rawMessagesById = new Map() - const rawIndexById = new Map() - for (const msg of rawMessages) { - rawMessagesById.set(msg.info.id, msg) - } - for (let index = 0; index < rawMessages.length; index++) { - const message = rawMessages[index] - if (!message) { - continue - } - rawIndexById.set(message.info.id, index) - } - - const summaryByBlockId = new Map() - for (const [blockId, block] of state.prune.messages.blocksById) { - if (!block.active) { - continue - } - summaryByBlockId.set(blockId, block) - } - - return { - rawMessages, - rawMessagesById, - rawIndexById, - summaryByBlockId, - } -} - -export function resolveBoundaryIds( - context: SearchContext, - state: SessionState, - startId: string, - endId: string, -): { startReference: BoundaryReference; endReference: BoundaryReference } { - const lookup = buildBoundaryReferenceLookup(context, state) - const issues: string[] = [] - const parsedStartId = parseBoundaryId(startId) - const parsedEndId = parseBoundaryId(endId) - - if (parsedStartId === null) { - issues.push("startId is invalid. Use an injected message ID (mNNNN) or block ID (bN).") - } - - if (parsedEndId === null) { - issues.push("endId is invalid. Use an injected message ID (mNNNN) or block ID (bN).") - } - - if (issues.length > 0) { - throwCombinedIssues(issues) - } - - if (!parsedStartId || !parsedEndId) { - throw new Error("Invalid boundary ID(s)") - } - - const startReference = lookup.get(parsedStartId.ref) - const endReference = lookup.get(parsedEndId.ref) - - if (!startReference) { - issues.push( - `startId ${parsedStartId.ref} is not available in the current conversation context. Choose an injected ID visible in context.`, - ) - } - - if (!endReference) { - issues.push( - `endId ${parsedEndId.ref} is not available in the current conversation context. Choose an injected ID visible in context.`, - ) - } - - if (issues.length > 0) { - throwCombinedIssues(issues) - } - - if (!startReference || !endReference) { - throw new Error("Failed to resolve boundary IDs") - } - - if (startReference.rawIndex > endReference.rawIndex) { - throw new Error( - `startId ${parsedStartId.ref} appears after endId ${parsedEndId.ref} in the conversation. Start must come before end.`, - ) - } - - return { startReference, endReference } -} - -function buildBoundaryReferenceLookup( - context: SearchContext, - state: SessionState, -): Map { - const lookup = new Map() - - for (const [messageRef, messageId] of state.messageIds.byRef) { - const rawMessage = context.rawMessagesById.get(messageId) - if (!rawMessage) { - continue - } - if (rawMessage.info.role === "user" && isIgnoredUserMessage(rawMessage)) { - continue - } - - const rawIndex = context.rawIndexById.get(messageId) - if (rawIndex === undefined) { - continue - } - lookup.set(messageRef, { - kind: "message", - rawIndex, - messageId, - }) - } - - const summaries = Array.from(context.summaryByBlockId.values()).sort( - (a, b) => a.blockId - b.blockId, - ) - for (const summary of summaries) { - const anchorMessage = context.rawMessagesById.get(summary.anchorMessageId) - if (!anchorMessage) { - continue - } - if (anchorMessage.info.role === "user" && isIgnoredUserMessage(anchorMessage)) { - continue - } - - const rawIndex = context.rawIndexById.get(summary.anchorMessageId) - if (rawIndex === undefined) { - continue - } - const blockRef = formatBlockRef(summary.blockId) - if (!lookup.has(blockRef)) { - lookup.set(blockRef, { - kind: "compressed-block", - rawIndex, - blockId: summary.blockId, - anchorMessageId: summary.anchorMessageId, - }) - } - } - - return lookup -} - -export function resolveRange( - context: SearchContext, - startReference: BoundaryReference, - endReference: BoundaryReference, -): RangeResolution { - const startRawIndex = startReference.rawIndex - const endRawIndex = endReference.rawIndex - const messageIds: string[] = [] - const messageSeen = new Set() - const toolIds: string[] = [] - const toolSeen = new Set() - const requiredBlockIds: number[] = [] - const requiredBlockSeen = new Set() - const messageTokenById = new Map() - - for (let index = startRawIndex; index <= endRawIndex; index++) { - const rawMessage = context.rawMessages[index] - if (!rawMessage) { - continue - } - if (rawMessage.info.role === "user" && isIgnoredUserMessage(rawMessage)) { - continue - } - - const messageId = rawMessage.info.id - if (!messageSeen.has(messageId)) { - messageSeen.add(messageId) - messageIds.push(messageId) - } - - if (!messageTokenById.has(messageId)) { - messageTokenById.set(messageId, countAllMessageTokens(rawMessage)) - } - - const parts = Array.isArray(rawMessage.parts) ? rawMessage.parts : [] - for (const part of parts) { - if (part.type !== "tool" || !part.callID) { - continue - } - if (toolSeen.has(part.callID)) { - continue - } - toolSeen.add(part.callID) - toolIds.push(part.callID) - } - } - - const rangeMessageIdSet = new Set(messageIds) - const summariesInRange: Array<{ blockId: number; rawIndex: number }> = [] - for (const summary of context.summaryByBlockId.values()) { - if (!rangeMessageIdSet.has(summary.anchorMessageId)) { - continue - } - - const anchorIndex = context.rawIndexById.get(summary.anchorMessageId) - if (anchorIndex === undefined) { - continue - } - - summariesInRange.push({ - blockId: summary.blockId, - rawIndex: anchorIndex, - }) - } - - summariesInRange.sort((a, b) => a.rawIndex - b.rawIndex || a.blockId - b.blockId) - for (const summary of summariesInRange) { - if (requiredBlockSeen.has(summary.blockId)) { - continue - } - requiredBlockSeen.add(summary.blockId) - requiredBlockIds.push(summary.blockId) - } - - if (messageIds.length === 0) { - throw new Error( - "Failed to map boundary matches back to raw messages. Choose boundaries that include original conversation messages.", - ) - } - - return { - startReference, - endReference, - messageIds, - messageTokenById, - toolIds, - requiredBlockIds, - } -} - -export function resolveAnchorMessageId(startReference: BoundaryReference): string { - if (startReference.kind === "compressed-block") { - if (!startReference.anchorMessageId) { - throw new Error("Failed to map boundary matches back to raw messages") - } - return startReference.anchorMessageId - } - - if (!startReference.messageId) { - throw new Error("Failed to map boundary matches back to raw messages") - } - return startReference.messageId -} - -export function parseBlockPlaceholders(summary: string): ParsedBlockPlaceholder[] { - const placeholders: ParsedBlockPlaceholder[] = [] - const regex = new RegExp(BLOCK_PLACEHOLDER_REGEX) - - let match: RegExpExecArray | null - while ((match = regex.exec(summary)) !== null) { - const full = match[0] - const blockIdPart = match[1] || match[2] - const parsed = Number.parseInt(blockIdPart, 10) - if (!Number.isInteger(parsed)) { - continue - } - - placeholders.push({ - raw: full, - blockId: parsed, - startIndex: match.index, - endIndex: match.index + full.length, - }) - } - - return placeholders -} - -export function validateSummaryPlaceholders( - placeholders: ParsedBlockPlaceholder[], - requiredBlockIds: number[], - startReference: BoundaryReference, - endReference: BoundaryReference, - summaryByBlockId: Map, -): number[] { - const boundaryOptionalIds = new Set() - if (startReference.kind === "compressed-block") { - if (startReference.blockId === undefined) { - throw new Error("Failed to map boundary matches back to raw messages") - } - boundaryOptionalIds.add(startReference.blockId) - } - if (endReference.kind === "compressed-block") { - if (endReference.blockId === undefined) { - throw new Error("Failed to map boundary matches back to raw messages") - } - boundaryOptionalIds.add(endReference.blockId) - } - - const strictRequiredIds = requiredBlockIds.filter((id) => !boundaryOptionalIds.has(id)) - const requiredSet = new Set(requiredBlockIds) - const keptPlaceholderIds = new Set() - const validPlaceholders: ParsedBlockPlaceholder[] = [] - - for (const placeholder of placeholders) { - const isKnown = summaryByBlockId.has(placeholder.blockId) - const isRequired = requiredSet.has(placeholder.blockId) - const isDuplicate = keptPlaceholderIds.has(placeholder.blockId) - - if (isKnown && isRequired && !isDuplicate) { - validPlaceholders.push(placeholder) - keptPlaceholderIds.add(placeholder.blockId) - } - } - - placeholders.length = 0 - placeholders.push(...validPlaceholders) - - return strictRequiredIds.filter((id) => !keptPlaceholderIds.has(id)) -} - -export function injectBlockPlaceholders( - summary: string, - placeholders: ParsedBlockPlaceholder[], - summaryByBlockId: Map, - startReference: BoundaryReference, - endReference: BoundaryReference, -): InjectedSummaryResult { - let cursor = 0 - let expanded = summary - const consumed: number[] = [] - const consumedSeen = new Set() - - if (placeholders.length > 0) { - expanded = "" - for (const placeholder of placeholders) { - const target = summaryByBlockId.get(placeholder.blockId) - if (!target) { - throw new Error( - `Compressed block not found: ${formatBlockPlaceholder(placeholder.blockId)}`, - ) - } - - expanded += summary.slice(cursor, placeholder.startIndex) - expanded += restoreSummary(target.summary) - cursor = placeholder.endIndex - - if (!consumedSeen.has(placeholder.blockId)) { - consumedSeen.add(placeholder.blockId) - consumed.push(placeholder.blockId) - } - } - - expanded += summary.slice(cursor) - } - - expanded = injectBoundarySummaryIfMissing( - expanded, - startReference, - "start", - summaryByBlockId, - consumed, - consumedSeen, - ) - expanded = injectBoundarySummaryIfMissing( - expanded, - endReference, - "end", - summaryByBlockId, - consumed, - consumedSeen, - ) - - return { - expandedSummary: expanded, - consumedBlockIds: consumed, - } -} - -export function allocateBlockId(state: SessionState): number { - const next = state.prune.messages.nextBlockId - if (!Number.isInteger(next) || next < 1) { - state.prune.messages.nextBlockId = 2 - return 1 - } - - state.prune.messages.nextBlockId = next + 1 - return next -} - -export function wrapCompressedSummary(blockId: number, summary: string): string { - const header = COMPRESSED_BLOCK_HEADER - const footer = formatMessageIdTag(formatBlockRef(blockId)) - const body = summary.trim() - if (body.length === 0) { - return `${header}\n${footer}` - } - return `${header}\n${body}\n\n${footer}` -} - -export function applyCompressionState( - state: SessionState, - input: CompressionStateInput, - range: RangeResolution, - anchorMessageId: string, - blockId: number, - summary: string, - consumedBlockIds: number[], -): AppliedCompressionResult { - const messagesState = state.prune.messages - const consumed = [...new Set(consumedBlockIds.filter((id) => Number.isInteger(id) && id > 0))] - const included = [...consumed] - - const effectiveMessageIds = new Set(range.messageIds) - const effectiveToolIds = new Set(range.toolIds) - - for (const consumedBlockId of consumed) { - const consumedBlock = messagesState.blocksById.get(consumedBlockId) - if (!consumedBlock) { - continue - } - for (const messageId of consumedBlock.effectiveMessageIds) { - effectiveMessageIds.add(messageId) - } - for (const toolId of consumedBlock.effectiveToolIds) { - effectiveToolIds.add(toolId) - } - } - - const initiallyActiveMessages = new Set() - for (const messageId of effectiveMessageIds) { - const entry = messagesState.byMessageId.get(messageId) - if (entry && entry.activeBlockIds.length > 0) { - initiallyActiveMessages.add(messageId) - } - } - - const initiallyActiveToolIds = new Set() - for (const activeBlockId of messagesState.activeBlockIds) { - const activeBlock = messagesState.blocksById.get(activeBlockId) - if (!activeBlock || !activeBlock.active) { - continue - } - - for (const toolId of activeBlock.effectiveToolIds) { - initiallyActiveToolIds.add(toolId) - } - } - - const createdAt = Date.now() - const block: CompressionBlock = { - blockId, - active: true, - deactivatedByUser: false, - compressedTokens: 0, - topic: input.topic, - startId: input.startId, - endId: input.endId, - anchorMessageId, - compressMessageId: input.compressMessageId, - includedBlockIds: included, - consumedBlockIds: consumed, - parentBlockIds: [], - directMessageIds: [], - directToolIds: [], - effectiveMessageIds: [...effectiveMessageIds], - effectiveToolIds: [...effectiveToolIds], - createdAt, - summary, - } - - messagesState.blocksById.set(blockId, block) - messagesState.activeBlockIds.add(blockId) - messagesState.activeByAnchorMessageId.set(anchorMessageId, blockId) - - const deactivatedAt = Date.now() - for (const consumedBlockId of consumed) { - const consumedBlock = messagesState.blocksById.get(consumedBlockId) - if (!consumedBlock || !consumedBlock.active) { - continue - } - - consumedBlock.active = false - consumedBlock.deactivatedAt = deactivatedAt - consumedBlock.deactivatedByBlockId = blockId - if (!consumedBlock.parentBlockIds.includes(blockId)) { - consumedBlock.parentBlockIds.push(blockId) - } - - messagesState.activeBlockIds.delete(consumedBlockId) - const mappedBlockId = messagesState.activeByAnchorMessageId.get( - consumedBlock.anchorMessageId, - ) - if (mappedBlockId === consumedBlockId) { - messagesState.activeByAnchorMessageId.delete(consumedBlock.anchorMessageId) - } - } - - const removeActiveBlockId = ( - entry: { activeBlockIds: number[] }, - blockIdToRemove: number, - ): void => { - if (entry.activeBlockIds.length === 0) { - return - } - entry.activeBlockIds = entry.activeBlockIds.filter((id) => id !== blockIdToRemove) - } - - for (const consumedBlockId of consumed) { - const consumedBlock = messagesState.blocksById.get(consumedBlockId) - if (!consumedBlock) { - continue - } - for (const messageId of consumedBlock.effectiveMessageIds) { - const entry = messagesState.byMessageId.get(messageId) - if (!entry) { - continue - } - removeActiveBlockId(entry, consumedBlockId) - } - } - - for (const messageId of range.messageIds) { - const tokenCount = range.messageTokenById.get(messageId) || 0 - const existing = messagesState.byMessageId.get(messageId) - - if (!existing) { - messagesState.byMessageId.set(messageId, { - tokenCount, - allBlockIds: [blockId], - activeBlockIds: [blockId], - }) - continue - } - - existing.tokenCount = Math.max(existing.tokenCount, tokenCount) - if (!existing.allBlockIds.includes(blockId)) { - existing.allBlockIds.push(blockId) - } - if (!existing.activeBlockIds.includes(blockId)) { - existing.activeBlockIds.push(blockId) - } - } - - for (const messageId of block.effectiveMessageIds) { - if (range.messageTokenById.has(messageId)) { - continue - } - - const existing = messagesState.byMessageId.get(messageId) - if (!existing) { - continue - } - if (!existing.allBlockIds.includes(blockId)) { - existing.allBlockIds.push(blockId) - } - if (!existing.activeBlockIds.includes(blockId)) { - existing.activeBlockIds.push(blockId) - } - } - - let compressedTokens = 0 - const newlyCompressedMessageIds: string[] = [] - for (const messageId of effectiveMessageIds) { - const entry = messagesState.byMessageId.get(messageId) - if (!entry) { - continue - } - - const isNowActive = entry.activeBlockIds.length > 0 - const wasActive = initiallyActiveMessages.has(messageId) - - if (isNowActive && !wasActive) { - compressedTokens += entry.tokenCount - newlyCompressedMessageIds.push(messageId) - } - } - - const newlyCompressedToolIds: string[] = [] - for (const toolId of effectiveToolIds) { - if (!initiallyActiveToolIds.has(toolId)) { - newlyCompressedToolIds.push(toolId) - } - } - - block.directMessageIds = [...newlyCompressedMessageIds] - block.directToolIds = [...newlyCompressedToolIds] - - block.compressedTokens = compressedTokens - - state.stats.pruneTokenCounter += compressedTokens - state.stats.totalPruneTokens += state.stats.pruneTokenCounter - state.stats.pruneTokenCounter = 0 - - return { - compressedTokens, - messageIds: range.messageIds, - newlyCompressedMessageIds, - newlyCompressedToolIds, - } -} - -function restoreSummary(summary: string): string { - const headerMatch = summary.match(/^\s*\[Compressed conversation(?: section)?(?: b\d+)?\]/i) - if (!headerMatch) { - return summary - } - - const afterHeader = summary.slice(headerMatch[0].length) - const withoutLeadingBreaks = afterHeader.replace(/^(?:\r?\n)+/, "") - return withoutLeadingBreaks - .replace(/(?:\r?\n)*b\d+<\/dcp-message-id>\s*$/i, "") - .replace(/(?:\r?\n)+$/, "") -} - -function injectBoundarySummaryIfMissing( - summary: string, - reference: BoundaryReference, - position: "start" | "end", - summaryByBlockId: Map, - consumed: number[], - consumedSeen: Set, -): string { - if (reference.kind !== "compressed-block" || reference.blockId === undefined) { - return summary - } - if (consumedSeen.has(reference.blockId)) { - return summary - } - - const target = summaryByBlockId.get(reference.blockId) - if (!target) { - throw new Error(`Compressed block not found: ${formatBlockPlaceholder(reference.blockId)}`) - } - - const injectedBody = restoreSummary(target.summary) - const next = - position === "start" - ? mergeWithSpacing(injectedBody, summary) - : mergeWithSpacing(summary, injectedBody) - - consumedSeen.add(reference.blockId) - consumed.push(reference.blockId) - return next -} - -function mergeWithSpacing(left: string, right: string): string { - const l = left.trim() - const r = right.trim() - - if (!l) { - return right - } - if (!r) { - return left - } - return `${l}\n\n${r}` -} - -export function appendProtectedUserMessages( - summary: string, - range: RangeResolution, - searchContext: SearchContext, - state: SessionState, - enabled: boolean, -): string { - if (!enabled) return summary - - const userTexts: string[] = [] - - for (const messageId of range.messageIds) { - const existingCompressionEntry = state.prune.messages.byMessageId.get(messageId) - if (existingCompressionEntry && existingCompressionEntry.activeBlockIds.length > 0) { - continue - } - - const message = searchContext.rawMessagesById.get(messageId) - if (!message) continue - if (message.info.role !== "user") continue - if (isIgnoredUserMessage(message)) continue - - const parts = Array.isArray(message.parts) ? message.parts : [] - for (const part of parts) { - if (part.type === "text" && typeof part.text === "string" && part.text.trim()) { - userTexts.push(part.text) - break - } - } - } - - if (userTexts.length === 0) { - return summary - } - - const heading = "\n\nThe following user messages were sent in this conversation verbatim:" - const body = userTexts.map((text) => `\n${text}`).join("") - return summary + heading + body -} - -export async function appendProtectedTools( - client: any, - state: SessionState, - allowSubAgents: boolean, - summary: string, - range: RangeResolution, - searchContext: SearchContext, - protectedTools: string[], - protectedFilePatterns: string[] = [], -): Promise { - const protectedOutputs: string[] = [] - - for (const messageId of range.messageIds) { - const existingCompressionEntry = state.prune.messages.byMessageId.get(messageId) - if (existingCompressionEntry && existingCompressionEntry.activeBlockIds.length > 0) { - continue - } - - const message = searchContext.rawMessagesById.get(messageId) - if (!message) continue - - const parts = Array.isArray(message.parts) ? message.parts : [] - for (const part of parts) { - if (part.type === "tool" && part.callID) { - let isToolProtected = isToolNameProtected(part.tool, protectedTools) - - if (!isToolProtected && protectedFilePatterns.length > 0) { - const filePaths = getFilePathsFromParameters(part.tool, part.state?.input) - if (isFilePathProtected(filePaths, protectedFilePatterns)) { - isToolProtected = true - } - } - - if (isToolProtected) { - const title = `Tool: ${part.tool}` - let output = "" - - if (part.state?.status === "completed" && part.state?.output) { - output = - typeof part.state.output === "string" - ? part.state.output - : JSON.stringify(part.state.output) - } - - if ( - allowSubAgents && - part.tool === "task" && - part.state?.status === "completed" && - typeof part.state?.output === "string" - ) { - const cachedSubAgentResult = state.subAgentResultCache.get(part.callID) - - if (cachedSubAgentResult !== undefined) { - if (cachedSubAgentResult) { - output = mergeSubagentResult( - part.state.output, - cachedSubAgentResult, - ) - } - } else { - const subAgentSessionId = getSubAgentId(part) - if (subAgentSessionId) { - let subAgentResultText = "" - try { - const subAgentMessages = await fetchSessionMessages( - client, - subAgentSessionId, - ) - subAgentResultText = buildSubagentResultText(subAgentMessages) - } catch { - subAgentResultText = "" - } - - if (subAgentResultText) { - state.subAgentResultCache.set(part.callID, subAgentResultText) - output = mergeSubagentResult( - part.state.output, - subAgentResultText, - ) - } - } - } - } - - if (output) { - protectedOutputs.push(`\n### ${title}\n${output}`) - } - } - } - } - } - - if (protectedOutputs.length === 0) { - return summary - } - - const heading = "\n\nThe following protected tools were used in this conversation as well:" - return summary + heading + protectedOutputs.join("") -} - -export function appendMissingBlockSummaries( - summary: string, - missingBlockIds: number[], - summaryByBlockId: Map, - consumedBlockIds: number[], -): InjectedSummaryResult { - const consumedSeen = new Set(consumedBlockIds) - const consumed = [...consumedBlockIds] - - const missingSummaries: string[] = [] - for (const blockId of missingBlockIds) { - if (consumedSeen.has(blockId)) { - continue - } - - const target = summaryByBlockId.get(blockId) - if (!target) { - throw new Error(`Compressed block not found: ${formatBlockPlaceholder(blockId)}`) - } - - missingSummaries.push( - `\n### ${formatBlockPlaceholder(blockId)}\n${restoreSummary(target.summary)}`, - ) - consumedSeen.add(blockId) - consumed.push(blockId) - } - - if (missingSummaries.length === 0) { - return { - expandedSummary: summary, - consumedBlockIds: consumed, - } - } - - const heading = - "\n\nThe following previously compressed summaries were also part of this conversation section:" - - return { - expandedSummary: summary + heading + missingSummaries.join(""), - consumedBlockIds: consumed, - } -} - -function throwCombinedIssues(issues: string[]): never { - if (issues.length === 1) { - throw new Error(issues[0]) - } - - throw new Error(issues.map((issue) => `- ${issue}`).join("\n")) -} diff --git a/lib/ui/notification.ts b/lib/ui/notification.ts index 80e766c9..e6909c64 100644 --- a/lib/ui/notification.ts +++ b/lib/ui/notification.ts @@ -16,6 +16,13 @@ export const PRUNE_REASON_LABELS: Record = { extraction: "Extraction", } +interface CompressionNotificationEntry { + blockId: number + runId: number + summary: string + summaryTokens: number +} + function buildMinimalMessage(state: SessionState, reason: PruneReason | undefined): string { const reasonSuffix = reason ? ` — ${PRUNE_REASON_LABELS[reason]}` : "" return ( @@ -127,15 +134,40 @@ export async function sendUnifiedNotification( return true } +function buildCompressionSummary( + entries: CompressionNotificationEntry[], + state: SessionState, +): string { + if (entries.length === 1) { + return entries[0]?.summary ?? "" + } + + return entries + .map((entry) => { + const topic = + state.prune.messages.blocksById.get(entry.blockId)?.topic ?? "(unknown topic)" + return `### ${topic}\n${entry.summary}` + }) + .join("\n\n") +} + +function getCompressionLabel(entries: CompressionNotificationEntry[]): string { + const runId = entries[0]?.runId + if (runId === undefined) { + return "Compression" + } + + return `Compression #${runId}` +} + export async function sendCompressNotification( client: any, logger: Logger, config: PluginConfig, state: SessionState, sessionId: string, - compressionId: number, - summary: string, - summaryTokens: number, + entries: CompressionNotificationEntry[], + batchTopic: string | undefined, totalSessionTokens: number, sessionMessageIds: string[], params: any, @@ -144,25 +176,66 @@ export async function sendCompressNotification( return false } + if (entries.length === 0) { + return false + } + let message: string + const compressionLabel = getCompressionLabel(entries) + const summary = buildCompressionSummary(entries, state) + const summaryTokens = entries.reduce((total, entry) => total + entry.summaryTokens, 0) const summaryTokensStr = formatTokenCount(summaryTokens) - const compressionBlock = state.prune.messages.blocksById.get(compressionId) + const compressedTokens = entries.reduce((total, entry) => { + const compressionBlock = state.prune.messages.blocksById.get(entry.blockId) + if (!compressionBlock) { + logger.error("Compression block missing for notification", { + compressionId: entry.blockId, + sessionId, + }) + return total + } - if (!compressionBlock) { - logger.error("Compression block missing for notification", { - compressionId, - sessionId, - }) + return total + compressionBlock.compressedTokens + }, 0) + + const newlyCompressedMessageIds: string[] = [] + const newlyCompressedToolIds: string[] = [] + const seenMessageIds = new Set() + const seenToolIds = new Set() + + for (const entry of entries) { + const compressionBlock = state.prune.messages.blocksById.get(entry.blockId) + if (!compressionBlock) { + continue + } + + for (const messageId of compressionBlock.directMessageIds) { + if (seenMessageIds.has(messageId)) { + continue + } + seenMessageIds.add(messageId) + newlyCompressedMessageIds.push(messageId) + } + + for (const toolId of compressionBlock.directToolIds) { + if (seenToolIds.has(toolId)) { + continue + } + seenToolIds.add(toolId) + newlyCompressedToolIds.push(toolId) + } } - const newlyCompressedToolIds = compressionBlock?.directToolIds ?? [] - const newlyCompressedMessageIds = compressionBlock?.directMessageIds ?? [] - const topic = compressionBlock?.topic ?? "(unknown topic)" - const compressedTokens = compressionBlock?.compressedTokens ?? 0 + const topic = + batchTopic ?? + (entries.length === 1 + ? (state.prune.messages.blocksById.get(entries[0]?.blockId ?? -1)?.topic ?? + "(unknown topic)") + : "(unknown topic)") if (config.pruneNotification === "minimal") { message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) - message += ` — Compression #${compressionId}` + message += ` — ${compressionLabel}` } else { message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) @@ -183,7 +256,7 @@ export async function sendCompressNotification( totalSessionTokens > 0 ? Math.round((compressedTokens / totalSessionTokens) * 100) : 0 message += `\n\n${progressBar}` - message += `\n▣ Compression #${compressionId} (${pruneTokenCounterStr} removed, ${reduction}% reduction)` + message += `\n▣ ${compressionLabel} (${pruneTokenCounterStr} removed, ${reduction}% reduction)` message += `\n→ Topic: ${topic}` message += `\n→ Items: ${newlyCompressedMessageIds.length} messages` if (newlyCompressedToolIds.length > 0) { diff --git a/scripts/opencode-message-token-counts b/scripts/opencode-message-token-counts new file mode 100755 index 00000000..b04022f5 --- /dev/null +++ b/scripts/opencode-message-token-counts @@ -0,0 +1,393 @@ +#!/usr/bin/env python3 +""" +Show countAllMessageTokens-style token counts for each message in an OpenCode session. + +Usage: opencode-message-token-counts [--session ID] [--json] [--no-color] [--db PATH] +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +from pathlib import Path +from typing import Optional + +from opencode_api import APIError, add_api_arguments, create_client_from_args, list_sessions_across_projects + + +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parent + + +class Colors: + RESET = "\033[0m" + BOLD = "\033[1m" + DIM = "\033[2m" + GREEN = "\033[32m" + YELLOW = "\033[33m" + CYAN = "\033[36m" + + +NO_COLOR = Colors() +for attr in dir(NO_COLOR): + if not attr.startswith("_"): + setattr(NO_COLOR, attr, "") + + +def stringify_json(value) -> str: + return json.dumps(value, separators=(",", ":"), ensure_ascii=False) + + +def collapse_whitespace(text: str) -> str: + return " ".join(text.split()) + + +def truncate(text: str, limit: int = 64) -> str: + if len(text) <= limit: + return text + return text[: limit - 3] + "..." + + +def get_terminal_width(default: int = 120) -> int: + return max(80, shutil.get_terminal_size((default, 20)).columns) + + +def short_message_id(message_id: str, limit: int = 14) -> str: + return truncate(message_id or "-", limit) + + +def preview_message(parts: list[dict]) -> str: + for part in parts: + if part.get("type") != "text": + continue + text = collapse_whitespace(part.get("text", "")) + if not text: + continue + prefix = "[ignored] " if part.get("ignored", False) else "" + return truncate(prefix + text) + + tool_names = [part.get("tool", "tool") for part in parts if part.get("type") == "tool"] + if tool_names: + return truncate(f"[tools: {', '.join(tool_names[:3])}]") + + for part in parts: + part_type = part.get("type", "unknown") + if part_type in {"step-start", "step-finish"}: + continue + if part_type == "tool": + tool_name = part.get("tool", "tool") + status = (part.get("state") or {}).get("status") + suffix = f" {status}" if status else "" + return f"[tool:{tool_name}{suffix}]" + return f"[{part_type}]" + + for part in parts: + part_type = part.get("type", "unknown") + return f"[{part_type}]" + + return "[no content]" + + +def extract_tool_content(part: dict) -> list[str]: + contents: list[str] = [] + tool_name = part.get("tool") + state = part.get("state") or {} + + if tool_name == "question": + questions = (state.get("input") or {}).get("questions") + if questions is not None: + content = questions if isinstance(questions, str) else stringify_json(questions) + contents.append(content) + return contents + + if tool_name in {"edit", "write"}: + if state.get("input") is not None: + input_content = state["input"] if isinstance(state["input"], str) else stringify_json(state["input"]) + contents.append(input_content) + + if state.get("status") == "completed" and state.get("output") is not None: + output = state["output"] + contents.append(output if isinstance(output, str) else stringify_json(output)) + elif state.get("status") == "error" and state.get("error") is not None: + error = state["error"] + contents.append(error if isinstance(error, str) else stringify_json(error)) + + return contents + + +def collect_message_segments(message: dict) -> tuple[list[str], int, int, list[str]]: + segments: list[str] = [] + text_segments = 0 + tool_segments = 0 + part_types: list[str] = [] + + for part in message.get("parts", []): + part_type = part.get("type", "unknown") + part_types.append(part_type) + if part_type == "text": + text = part.get("text", "") + if text: + segments.append(text) + text_segments += 1 + continue + + tool_contents = extract_tool_content(part) + segments.extend(tool_contents) + tool_segments += len(tool_contents) + + return segments, text_segments, tool_segments, part_types + + +def fallback_count_tokens(text: str) -> int: + if not text: + return 0 + return round(len(text) / 4) + + +def count_tokens_batch(texts: list[str]) -> tuple[list[int], str]: + if not texts: + return [], "anthropic" + + node_script = """ +import { countTokens } from \"@anthropic-ai/tokenizer\"; +import { readFileSync } from \"node:fs\"; + +const texts = JSON.parse(readFileSync(0, \"utf8\")); +const counts = texts.map((text) => countTokens(text || \"\")); +process.stdout.write(JSON.stringify(counts)); +""".strip() + + try: + proc = subprocess.run( + ["node", "--input-type=module", "-e", node_script], + input=stringify_json(texts), + capture_output=True, + text=True, + cwd=REPO_ROOT, + check=True, + timeout=15, + ) + counts = json.loads(proc.stdout) + if isinstance(counts, list) and len(counts) == len(texts): + return [int(count) for count in counts], "anthropic" + except (subprocess.SubprocessError, FileNotFoundError, json.JSONDecodeError, ValueError): + pass + + return [fallback_count_tokens(text) for text in texts], "approximate" + + +def get_most_recent_session(client, session_list_limit: int) -> Optional[dict]: + sessions = list_sessions_across_projects(client, per_project_limit=session_list_limit) + return sessions[0] if sessions else None + + +def analyze_session(client, session: dict) -> dict: + session_id = session["id"] + messages = client.get_session_messages(session_id, directory=session.get("directory")) + + analyzed_messages = [] + count_inputs: list[str] = [] + for index, message in enumerate(messages, 1): + info = message.get("info", {}) + segments, text_segments, tool_segments, part_types = collect_message_segments(message) + count_inputs.append(" ".join(segments)) + analyzed_messages.append( + { + "index": index, + "message_id": info.get("id", ""), + "role": info.get("role", "unknown"), + "part_count": len(message.get("parts", [])), + "part_types": part_types, + "counted_segments": len(segments), + "text_segments": text_segments, + "tool_segments": tool_segments, + "preview": preview_message(message.get("parts", [])), + } + ) + + counts, tokenizer = count_tokens_batch(count_inputs) + total_tokens = 0 + nonzero_messages = 0 + max_tokens = 0 + + for message_data, count in zip(analyzed_messages, counts): + message_data["tokens"] = count + total_tokens += count + if count > 0: + nonzero_messages += 1 + max_tokens = max(max_tokens, count) + + return { + "session_id": session_id, + "title": session.get("title", "Unknown"), + "tokenizer": tokenizer, + "messages": analyzed_messages, + "total_messages": len(analyzed_messages), + "messages_with_tokens": nonzero_messages, + "messages_without_tokens": len(analyzed_messages) - nonzero_messages, + "total_tokens": total_tokens, + "max_message_tokens": max_tokens, + } + + +def format_token_count(count: int, colors: Colors) -> str: + c = colors + if count == 0: + return f"{c.DIM}{count:>10,}{c.RESET}" + return f"{count:>10,}" + + +def format_role(role: str, colors: Colors, width: int = 9) -> str: + c = colors + label = f"{role:<{width}}" + if role == "user": + return f"{c.CYAN}{label}{c.RESET}" + if role == "assistant": + return f"{c.GREEN}{label}{c.RESET}" + return f"{c.YELLOW}{label}{c.RESET}" + + +def format_size_indicator(count: int, max_count: int, width: int = 8) -> str: + if max_count <= 0: + return f"{'.' * width} 0%" + + pct = round((count / max_count) * 100) + if count <= 0: + filled = 0 + else: + filled = max(1, round((count / max_count) * width)) + filled = min(width, filled) + return f"{'#' * filled}{'.' * (width - filled)} {pct:>3}%" + + +def largest_messages(messages: list[dict], limit: int = 5) -> list[dict]: + return sorted(messages, key=lambda message: message.get("tokens", 0), reverse=True)[:limit] + + +def print_wide_message_table(result: dict, colors: Colors, width: int): + c = colors + messages = result["messages"] + preview_width = max(24, width - 72) + + print( + f"{c.BOLD}{'#':>3} {'Role':<9} {'Tokens':>10} {'Size':<12} {'Seg/Part':<8} {'ID':<14} Preview{c.RESET}" + ) + print("-" * width) + + for message in messages: + preview = truncate(message["preview"], preview_width) + mix = f"{message['counted_segments']}/{message['part_count']}" + print( + f"{message['index']:>3} " + f"{format_role(message['role'], c, 9)} " + f"{format_token_count(message['tokens'], c)} " + f"{format_size_indicator(message['tokens'], result['max_message_tokens']):<12} " + f"{mix:<8} " + f"{c.DIM}{short_message_id(message['message_id']):<14}{c.RESET} " + f"{preview}" + ) + + +def print_compact_message_list(result: dict, colors: Colors, width: int): + c = colors + messages = result["messages"] + meta_width = max(18, width - 6) + preview_width = max(32, width - 8) + + print(f"{c.BOLD}Messages{c.RESET}") + print("-" * width) + + for message in messages: + tokens = f"{message['tokens']:,} tokens" + size = format_size_indicator(message["tokens"], result["max_message_tokens"]) + mix = f"{message['counted_segments']}/{message['part_count']} seg/part" + meta = truncate(f"{tokens} {size} {mix}", meta_width) + preview = truncate(message["preview"], preview_width) + + print(f"{message['index']:>3} {format_role(message['role'], c, 9)} {meta}") + print(f" {c.DIM}{short_message_id(message['message_id'])}{c.RESET} {preview}") + + +def print_highlights(result: dict, colors: Colors, width: int): + c = colors + heavy_messages = [message for message in largest_messages(result["messages"]) if message.get("tokens", 0) > 0] + if not heavy_messages: + return + + print(f"\n{c.BOLD}Largest messages{c.RESET}") + print("-" * width) + for message in heavy_messages: + print( + f" #{message['index']:<3} {format_role(message['role'], c, 9)} " + f"{message['tokens']:>10,} {truncate(message['preview'], max(30, width - 33))}" + ) + + +def print_message_tokens(result: dict, colors: Colors): + c = colors + width = get_terminal_width() + print(f"{c.BOLD}{'=' * width}{c.RESET}") + print(f"{c.BOLD}SESSION MESSAGE TOKEN COUNTS{c.RESET}") + print(f"{c.BOLD}{'=' * width}{c.RESET}\n") + print(f" Session: {c.CYAN}{result['session_id']}{c.RESET}") + print(f" Title: {result['title']}") + print(f" Messages: {result['total_messages']}") + print(f" Tokenizer: {result['tokenizer']}") + print(f" Total: {result['total_tokens']:,} tokens") + print(f" Largest: {result['max_message_tokens']:,} tokens\n") + + if not result["messages"]: + print(" No messages found in this session.") + return + + if width >= 110: + print_wide_message_table(result, c, width) + else: + print_compact_message_list(result, c, width) + + print("-" * width) + print_highlights(result, c, width) + print(f"\n{c.BOLD}SESSION SUMMARY{c.RESET}") + print(f" Total message tokens: {result['total_tokens']:,}") + print(f" Messages with tokens: {result['messages_with_tokens']:,}") + print(f" Empty messages: {result['messages_without_tokens']:,}") + print(f" Largest message: {result['max_message_tokens']:,}") + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Show countAllMessageTokens-style token counts for each message in an OpenCode session" + ) + parser.add_argument("--session", "-s", type=str, default=None, help="Session ID to analyze (default: most recent)") + parser.add_argument("--json", "-j", action="store_true", help="Output as JSON") + parser.add_argument("--no-color", action="store_true", help="Disable colored output") + add_api_arguments(parser) + args = parser.parse_args() + + try: + with create_client_from_args(args) as client: + if args.session is None: + session = get_most_recent_session(client, args.session_list_limit) + if session is None: + print("Error: No sessions found") + return 1 + else: + session = client.get_session(args.session) + result = analyze_session(client, session) + except APIError as err: + print(f"Error: {err}") + return 1 + + if args.json: + print(json.dumps(result, indent=2, ensure_ascii=False)) + else: + colors = NO_COLOR if args.no_color else Colors() + print_message_tokens(result, colors) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/print.ts b/scripts/print.ts index e8c0f5e9..9e5f4eea 100644 --- a/scripts/print.ts +++ b/scripts/print.ts @@ -13,8 +13,10 @@ function getPromptByKey(prompts: RuntimePrompts, key: PromptKey): string { switch (key) { case "system": return prompts.system - case "compress": - return prompts.compress + case "compress-range": + return prompts.compressRange + case "compress-message": + return prompts.compressMessage case "context-limit-nudge": return prompts.contextLimitNudge case "turn-nudge": @@ -43,12 +45,12 @@ Options: --system-all Print system prompt with both overlays Prompt keys: - system, compress, context-limit-nudge, - turn-nudge, iteration-nudge + system, compress-range, compress-message, + context-limit-nudge, turn-nudge, iteration-nudge Examples: npm run dcp -- --list - npm run dcp -- --show compress + npm run dcp -- --show compress-range npm run dcp -- --system-all `) process.exit(0) diff --git a/tests/compress-message.test.ts b/tests/compress-message.test.ts new file mode 100644 index 00000000..a48a2f9b --- /dev/null +++ b/tests/compress-message.test.ts @@ -0,0 +1,649 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { join } from "node:path" +import { tmpdir } from "node:os" +import { mkdirSync } from "node:fs" +import { createCompressMessageTool } from "../lib/compress/message" +import { createSessionState, type WithParts } from "../lib/state" +import type { PluginConfig } from "../lib/config" +import { Logger } from "../lib/logger" + +const testDataHome = join(tmpdir(), `opencode-dcp-message-tests-${process.pid}`) +const testConfigHome = join(tmpdir(), `opencode-dcp-message-config-tests-${process.pid}`) + +process.env.XDG_DATA_HOME = testDataHome +process.env.XDG_CONFIG_HOME = testConfigHome + +mkdirSync(testDataHome, { recursive: true }) +mkdirSync(testConfigHome, { recursive: true }) + +function buildConfig(): PluginConfig { + return { + enabled: true, + debug: false, + pruneNotification: "off", + pruneNotificationType: "chat", + commands: { + enabled: true, + protectedTools: [], + }, + manualMode: { + enabled: false, + automaticStrategies: true, + }, + turnProtection: { + enabled: false, + turns: 4, + }, + experimental: { + allowSubAgents: false, + customPrompts: false, + }, + protectedFilePatterns: [], + compress: { + mode: "message", + permission: "allow", + showCompression: false, + maxContextLimit: 150000, + minContextLimit: 50000, + nudgeFrequency: 5, + iterationNudgeThreshold: 15, + nudgeForce: "soft", + protectedTools: ["task"], + protectUserMessages: false, + }, + strategies: { + deduplication: { + enabled: true, + protectedTools: [], + }, + purgeErrors: { + enabled: true, + turns: 4, + protectedTools: [], + }, + }, + } +} + +function textPart(messageID: string, sessionID: string, id: string, text: string) { + return { + id, + messageID, + sessionID, + type: "text" as const, + text, + } +} + +function toolPart( + messageID: string, + sessionID: string, + callID: string, + toolName: string, + output: string, +) { + return { + id: `${callID}-part`, + messageID, + sessionID, + type: "tool" as const, + tool: toolName, + callID, + state: { + status: "completed" as const, + input: { description: "demo" }, + output, + }, + } +} + +function buildMessages(sessionID: string): WithParts[] { + return [ + { + info: { + id: "msg-user-1", + role: "user", + sessionID, + agent: "assistant", + model: { + providerID: "anthropic", + modelID: "claude-test", + }, + time: { created: 1 }, + } as WithParts["info"], + parts: [textPart("msg-user-1", sessionID, "part-1", "Investigate the issue")], + }, + { + info: { + id: "msg-assistant-1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 2 }, + } as WithParts["info"], + parts: [textPart("msg-assistant-1", sessionID, "part-2", "I mapped the code path")], + }, + { + info: { + id: "msg-assistant-2", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 3 }, + } as WithParts["info"], + parts: [ + textPart("msg-assistant-2", sessionID, "part-3", "I also ran a task tool"), + toolPart("msg-assistant-2", sessionID, "call-task-1", "task", "task output body"), + ], + }, + ] +} + +test("compress message tool appends non-editable format overlay", () => { + const tool = createCompressMessageTool({ + client: {}, + state: createSessionState(), + logger: new Logger(false), + config: buildConfig(), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + assert.match(tool.description, /THE FORMAT OF COMPRESS/) + assert.match(tool.description, /messageId: string/) + assert.match(tool.description, /Raw message ID only: mNNNN/) +}) + +test("compress message mode batches individual message summaries", async () => { + const sessionID = `ses_message_compress_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config: buildConfig(), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + const result = await tool.execute( + { + topic: "Batch stale notes", + content: [ + { + messageId: "m0002", + topic: "Code path note", + summary: "Captured the assistant's code-path findings.", + }, + { + messageId: "m0003", + topic: "Task output note", + summary: "Captured the assistant's task-backed follow-up.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message", + }, + ) + + assert.equal(result, "Compressed 2 messages into [Compressed conversation section].") + assert.equal(state.prune.messages.blocksById.size, 2) + + const blocks = Array.from(state.prune.messages.blocksById.values()).sort( + (a, b) => a.blockId - b.blockId, + ) + + assert.equal(blocks[0]?.startId, "m0002") + assert.equal(blocks[0]?.endId, "m0002") + assert.equal(blocks[0]?.topic, "Code path note") + assert.equal(blocks[1]?.startId, "m0003") + assert.equal(blocks[1]?.endId, "m0003") + assert.match( + blocks[1]?.summary || "", + /The following protected tools were used in this conversation as well:/, + ) + assert.match(blocks[1]?.summary || "", /Tool: task/) + assert.match(blocks[1]?.summary || "", /task output body/) +}) + +test("compress message mode does not partially apply when preparation fails", async () => { + const sessionID = `ses_message_compress_prepare_fail_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const config = buildConfig() + config.experimental.allowSubAgents = true + + state.subAgentResultCache.get = (() => { + throw new Error("cache failure") + }) as typeof state.subAgentResultCache.get + + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config, + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + await assert.rejects( + tool.execute( + { + topic: "Batch stale notes", + content: [ + { + messageId: "m0002", + topic: "Code path note", + summary: "Captured the assistant's code-path findings.", + }, + { + messageId: "m0003", + topic: "Task output note", + summary: "Captured the assistant's task-backed follow-up.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-prepare-fail", + }, + ), + /cache failure/, + ) + + assert.equal(state.prune.messages.blocksById.size, 0) +}) + +test("compress message mode rejects compressed block ids", async () => { + const sessionID = `ses_message_compress_reject_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config: buildConfig(), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + await assert.rejects( + tool.execute( + { + topic: "Reject block ids", + content: [ + { + messageId: "b1", + topic: "Invalid target", + summary: "Should not be accepted.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-reject", + }, + ), + /Unable to compress any messages\. Found 1 issue:/, + ) +}) + +test("compress message mode skips protected user message references", async () => { + const sessionID = `ses_message_compress_protected_user_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const config = buildConfig() + config.compress.protectUserMessages = true + + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config, + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + const result = await tool.execute( + { + topic: "Protected user entries", + content: [ + { + messageId: "BLOCKED", + topic: "Protected marker", + summary: "Should be skipped.", + }, + { + messageId: "m0001", + topic: "Hidden protected ref", + summary: "Should also be skipped.", + }, + { + messageId: "m0002", + topic: "Valid note", + summary: "Captured the assistant's code-path findings.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-protected-user", + }, + ) + + assert.equal(state.prune.messages.blocksById.size, 1) + assert.match(result, /^Compressed 1 message into \[Compressed conversation section\]\./) + assert.match(result, /Skipped 2 issues:/) + assert.match(result, /messageId BLOCKED refers to a protected message/) + assert.match(result, /messageId m0001 refers to a protected message/) +}) + +test("compress message mode allows messages containing compress tool parts", async () => { + const sessionID = `ses_message_compress_tool_${Date.now()}` + const rawMessages = buildMessages(sessionID) + rawMessages.push({ + info: { + id: "msg-assistant-compress", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 4 }, + } as WithParts["info"], + parts: [ + { + id: "compress-part", + messageID: "msg-assistant-compress", + sessionID, + type: "tool" as const, + tool: "compress", + callID: "call-compress-1", + state: { + status: "completed" as const, + input: { topic: "Earlier compression" }, + output: "done", + }, + }, + ], + }) + + const state = createSessionState() + const logger = new Logger(false) + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config: buildConfig(), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + const result = await tool.execute( + { + topic: "Compress compress call", + content: [ + { + messageId: "m0004", + topic: "Compress tool message", + summary: "Captured the earlier compress tool call.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-allow-compress-tool", + }, + ) + + assert.equal(result, "Compressed 1 message into [Compressed conversation section].") + assert.equal(state.prune.messages.blocksById.size, 1) + const block = Array.from(state.prune.messages.blocksById.values())[0] + assert.equal(block?.startId, "m0004") +}) + +test("compress message mode sends one aggregated notification for batched messages", async () => { + const sessionID = `ses_message_compress_notify_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const config = buildConfig() + config.pruneNotification = "detailed" + config.pruneNotificationType = "toast" + + const toastCalls: string[] = [] + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + tui: { + showToast: async ({ body }: { body: { message: string } }) => { + toastCalls.push(body.message) + }, + }, + }, + state, + logger, + config, + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + await tool.execute( + { + topic: "Batch stale notes", + content: [ + { + messageId: "m0002", + topic: "Code path note", + summary: "Captured the assistant's code-path findings.", + }, + { + messageId: "m0003", + topic: "Task output note", + summary: "Captured the assistant's task-backed follow-up.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-notify", + }, + ) + + assert.equal(toastCalls.length, 1) + assert.match(toastCalls[0] || "", /Compression #1/) + assert.match(toastCalls[0] || "", /Topic: Batch stale notes/) + assert.match(toastCalls[0] || "", /Items: 2 messages/) +}) + +test("compress message mode skips invalid batch entries and reports issues", async () => { + const sessionID = `ses_message_compress_partial_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config: buildConfig(), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + const result = await tool.execute( + { + topic: "Mixed entries", + content: [ + { + messageId: "b1", + topic: "Invalid block id", + summary: "Should be skipped.", + }, + { + messageId: "m0002", + topic: "Valid note", + summary: "Captured the assistant's code-path findings.", + }, + { + messageId: "m9999", + topic: "Missing message", + summary: "Should also be skipped.", + }, + { + messageId: "m0002", + topic: "Duplicate valid note", + summary: "Duplicate entry should be skipped.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-partial", + }, + ) + + assert.equal(state.prune.messages.blocksById.size, 1) + assert.match(result, /^Compressed 1 message into \[Compressed conversation section\]\./) + assert.match(result, /Skipped 3 issues:/) + assert.match(result, /Block IDs like bN are not allowed/) + assert.match(result, /messageId m9999 is not available in the current conversation context/) + assert.match(result, /messageId m0002 was selected more than once in this batch\./) +}) + +test("compress message mode reports issues when every batch entry is skipped", async () => { + const sessionID = `ses_message_compress_all_invalid_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const tool = createCompressMessageTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config: buildConfig(), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressMessage: "", compressRange: "" } + }, + }, + } as any) + + await assert.rejects( + tool.execute( + { + topic: "All invalid", + content: [ + { + messageId: "b1", + topic: "Invalid block id", + summary: "Should be skipped.", + }, + { + messageId: "m9999", + topic: "Missing message", + summary: "Should also be skipped.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-all-invalid", + }, + ), + /Unable to compress any messages\. Found 2 issues:/, + ) + + assert.equal(state.prune.messages.blocksById.size, 0) +}) diff --git a/tests/compress-placeholders.test.ts b/tests/compress-range-placeholders.test.ts similarity index 89% rename from tests/compress-placeholders.test.ts rename to tests/compress-range-placeholders.test.ts index e7b1c27a..e3a8a1ef 100644 --- a/tests/compress-placeholders.test.ts +++ b/tests/compress-range-placeholders.test.ts @@ -6,13 +6,14 @@ import { injectBlockPlaceholders, parseBlockPlaceholders, validateSummaryPlaceholders, - wrapCompressedSummary, - type BoundaryReference, -} from "../lib/tools/utils" +} from "../lib/compress/range-utils" +import { wrapCompressedSummary } from "../lib/compress/state" +import type { BoundaryReference } from "../lib/compress/types" function createBlock(blockId: number, body: string): CompressionBlock { return { blockId, + runId: blockId, active: true, deactivatedByUser: false, compressedTokens: 0, @@ -41,7 +42,7 @@ function createMessageBoundary(messageId: string, rawIndex: number): BoundaryRef } } -test("compress placeholder validation keeps valid placeholders and ignores invalid ones", () => { +test("compress range placeholder validation keeps valid placeholders and ignores invalid ones", () => { const summaryByBlockId = new Map([ [1, createBlock(1, "First compressed summary")], [2, createBlock(2, "Second compressed summary")], @@ -78,7 +79,7 @@ test("compress placeholder validation keeps valid placeholders and ignores inval assert.deepEqual(injected.consumedBlockIds, [1]) }) -test("compress continues by appending required block summaries the model omitted", () => { +test("compress range continues by appending required block summaries the model omitted", () => { const summaryByBlockId = new Map([[1, createBlock(1, "Recovered compressed summary")]]) const summary = "The model forgot to include the prior block." const parsed = parseBlockPlaceholders(summary) diff --git a/tests/compress.test.ts b/tests/compress-range.test.ts similarity index 53% rename from tests/compress.test.ts rename to tests/compress-range.test.ts index 988603e2..621c6622 100644 --- a/tests/compress.test.ts +++ b/tests/compress-range.test.ts @@ -3,7 +3,7 @@ import test from "node:test" import { join } from "node:path" import { tmpdir } from "node:os" import { mkdirSync } from "node:fs" -import { createCompressTool } from "../lib/tools/compress" +import { createCompressRangeTool } from "../lib/compress/range" import { createSessionState, type WithParts } from "../lib/state" import type { PluginConfig } from "../lib/config" import { Logger } from "../lib/logger" @@ -41,6 +41,7 @@ function buildConfig(): PluginConfig { }, protectedFilePatterns: [], compress: { + mode: "range", permission: "allow", showCompression: false, maxContextLimit: 150000, @@ -48,7 +49,6 @@ function buildConfig(): PluginConfig { nudgeFrequency: 5, iterationNudgeThreshold: 15, nudgeForce: "soft", - flatSchema: false, protectedTools: [], protectUserMessages: false, }, @@ -57,9 +57,6 @@ function buildConfig(): PluginConfig { enabled: true, protectedTools: [], }, - supersedeWrites: { - enabled: true, - }, purgeErrors: { enabled: true, turns: 4, @@ -126,7 +123,7 @@ function buildMessages(sessionID: string): WithParts[] { ] } -test("compress rebuilds subagent message refs after session state was reset", async () => { +test("compress range rebuilds subagent message refs after session state was reset", async () => { const sessionID = `ses_subagent_compress_${Date.now()}` const rawMessages = buildMessages(sessionID) const state = createSessionState() @@ -136,7 +133,7 @@ test("compress rebuilds subagent message refs after session state was reset", as state.messageIds.nextRef = 2 const logger = new Logger(false) - const tool = createCompressTool({ + const tool = createCompressRangeTool({ client: { session: { messages: async () => ({ data: rawMessages }), @@ -149,7 +146,7 @@ test("compress rebuilds subagent message refs after session state was reset", as prompts: { reload() {}, getRuntimePrompts() { - return { compress: "" } + return { compressRange: "", compressMessage: "" } }, }, } as any) @@ -157,11 +154,13 @@ test("compress rebuilds subagent message refs after session state was reset", as const result = await tool.execute( { topic: "Subagent race fix", - content: { - startId: "m0001", - endId: "m0002", - summary: "Captured the initial investigation and follow-up request.", - }, + content: [ + { + startId: "m0001", + endId: "m0002", + summary: "Captured the initial investigation and follow-up request.", + }, + ], }, { ask: async () => {}, @@ -178,3 +177,121 @@ test("compress rebuilds subagent message refs after session state was reset", as assert.equal(state.messageIds.byRef.get("m0002"), "msg-user-2") assert.equal(state.prune.messages.blocksById.size, 1) }) + +test("compress range mode batches multiple ranges into one notification", async () => { + const sessionID = `ses_range_compress_batch_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const config = buildConfig() + config.pruneNotification = "detailed" + config.pruneNotificationType = "toast" + + const toastCalls: string[] = [] + const tool = createCompressRangeTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: "ses_parent" } }), + }, + tui: { + showToast: async ({ body }: { body: { message: string } }) => { + toastCalls.push(body.message) + }, + }, + }, + state, + logger, + config, + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressRange: "", compressMessage: "" } + }, + }, + } as any) + + const result = await tool.execute( + { + topic: "Batch stale notes", + content: [ + { + startId: "m0001", + endId: "m0001", + summary: "Captured the initial assistant investigation.", + }, + { + startId: "m0002", + endId: "m0002", + summary: "Captured the follow-up user request.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-range-batch", + }, + ) + + assert.equal(result, "Compressed 2 messages into [Compressed conversation section].") + assert.equal(state.prune.messages.blocksById.size, 2) + assert.equal(toastCalls.length, 1) + assert.match(toastCalls[0] || "", /Compression #1/) + assert.match(toastCalls[0] || "", /Topic: Batch stale notes/) + assert.match(toastCalls[0] || "", /Items: 2 messages/) +}) + +test("compress range mode rejects overlapping batched ranges", async () => { + const sessionID = `ses_range_compress_overlap_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const tool = createCompressRangeTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: "ses_parent" } }), + }, + }, + state, + logger, + config: buildConfig(), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressRange: "", compressMessage: "" } + }, + }, + } as any) + + await assert.rejects( + tool.execute( + { + topic: "Overlapping ranges", + content: [ + { + startId: "m0001", + endId: "m0002", + summary: "Captured the initial investigation and follow-up request.", + }, + { + startId: "m0002", + endId: "m0002", + summary: "Captured the follow-up request again.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-range-overlap", + }, + ), + /Overlapping ranges cannot be compressed in the same batch/, + ) + + assert.equal(state.prune.messages.blocksById.size, 0) +}) diff --git a/tests/compression-groups.test.ts b/tests/compression-groups.test.ts new file mode 100644 index 00000000..f97d02ba --- /dev/null +++ b/tests/compression-groups.test.ts @@ -0,0 +1,447 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { join } from "node:path" +import { tmpdir } from "node:os" +import { mkdirSync } from "node:fs" +import { createCompressMessageTool } from "../lib/compress/message" +import { createCompressRangeTool } from "../lib/compress/range" +import { handleDecompressCommand } from "../lib/commands/decompress" +import { handleRecompressCommand } from "../lib/commands/recompress" +import { createSessionState, type WithParts } from "../lib/state" +import type { PluginConfig } from "../lib/config" +import { Logger } from "../lib/logger" + +const testDataHome = join(tmpdir(), `opencode-dcp-compression-groups-${process.pid}`) +const testConfigHome = join(tmpdir(), `opencode-dcp-compression-groups-config-${process.pid}`) + +process.env.XDG_DATA_HOME = testDataHome +process.env.XDG_CONFIG_HOME = testConfigHome + +mkdirSync(testDataHome, { recursive: true }) +mkdirSync(testConfigHome, { recursive: true }) + +function buildConfig(mode: "message" | "range"): PluginConfig { + return { + enabled: true, + debug: false, + pruneNotification: "off", + pruneNotificationType: "chat", + commands: { + enabled: true, + protectedTools: [], + }, + manualMode: { + enabled: false, + automaticStrategies: true, + }, + turnProtection: { + enabled: false, + turns: 4, + }, + experimental: { + allowSubAgents: false, + customPrompts: false, + }, + protectedFilePatterns: [], + compress: { + mode, + permission: "allow", + showCompression: false, + maxContextLimit: 150000, + minContextLimit: 50000, + nudgeFrequency: 5, + iterationNudgeThreshold: 15, + nudgeForce: "soft", + protectedTools: ["task"], + protectUserMessages: false, + }, + strategies: { + deduplication: { + enabled: true, + protectedTools: [], + }, + purgeErrors: { + enabled: true, + turns: 4, + protectedTools: [], + }, + }, + } +} + +function textPart(messageID: string, sessionID: string, id: string, text: string) { + return { + id, + messageID, + sessionID, + type: "text" as const, + text, + } +} + +function toolPart( + messageID: string, + sessionID: string, + callID: string, + toolName: string, + output: string, +) { + return { + id: `${callID}-part`, + messageID, + sessionID, + type: "tool" as const, + tool: toolName, + callID, + state: { + status: "completed" as const, + input: { description: "demo" }, + output, + }, + } +} + +function buildMessages(sessionID: string): WithParts[] { + return [ + { + info: { + id: "msg-user-1", + role: "user", + sessionID, + agent: "assistant", + model: { + providerID: "anthropic", + modelID: "claude-test", + }, + time: { created: 1 }, + } as WithParts["info"], + parts: [textPart("msg-user-1", sessionID, "part-1", "Investigate the issue")], + }, + { + info: { + id: "msg-assistant-1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 2 }, + } as WithParts["info"], + parts: [textPart("msg-assistant-1", sessionID, "part-2", "I mapped the code path")], + }, + { + info: { + id: "msg-assistant-2", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 3 }, + } as WithParts["info"], + parts: [ + textPart("msg-assistant-2", sessionID, "part-3", "I also ran a task tool"), + toolPart("msg-assistant-2", sessionID, "call-task-1", "task", "task output body"), + ], + }, + ] +} + +function appendOriginMessage(rawMessages: WithParts[], sessionID: string, messageID: string): void { + rawMessages.push({ + info: { + id: messageID, + role: "assistant", + sessionID, + agent: "assistant", + time: { created: rawMessages.length + 1 }, + } as WithParts["info"], + parts: [textPart(messageID, sessionID, `${messageID}-part`, "compress tool output")], + }) +} + +test("compression notifications increment by tool call across range and message tools", async () => { + const sessionID = `ses_compression_notifications_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const toastCalls: string[] = [] + const client = { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + tui: { + showToast: async ({ body }: { body: { message: string } }) => { + toastCalls.push(body.message) + }, + }, + } + + const rangeConfig = buildConfig("range") + rangeConfig.pruneNotification = "detailed" + rangeConfig.pruneNotificationType = "toast" + const messageConfig = buildConfig("message") + messageConfig.pruneNotification = "detailed" + messageConfig.pruneNotificationType = "toast" + + const rangeTool = createCompressRangeTool({ + client, + state, + logger, + config: rangeConfig, + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressRange: "", compressMessage: "" } + }, + }, + } as any) + + await rangeTool.execute( + { + topic: "Range batch", + content: [ + { + startId: "m0001", + endId: "m0001", + summary: "Captured the opening user request.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-range-origin", + }, + ) + + appendOriginMessage(rawMessages, sessionID, "msg-compress-range-origin") + + const messageTool = createCompressMessageTool({ + client, + state, + logger, + config: messageConfig, + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressRange: "", compressMessage: "" } + }, + }, + } as any) + + await messageTool.execute( + { + topic: "Message batch", + content: [ + { + messageId: "m0002", + topic: "Code path note", + summary: "Captured the assistant code-path findings.", + }, + { + messageId: "m0003", + topic: "Task output note", + summary: "Captured the assistant task-backed follow-up.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-origin", + }, + ) + + assert.equal(toastCalls.length, 2) + assert.match(toastCalls[0] || "", /Compression #1/) + assert.match(toastCalls[1] || "", /Compression #2/) +}) + +test("decompress groups batched message compressions by tool call", async () => { + const sessionID = `ses_message_grouped_decompress_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const ignoredMessages: string[] = [] + const client = { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + prompt: async ({ body }: { body: { parts: Array<{ text: string }> } }) => { + ignoredMessages.push(body.parts[0]?.text || "") + }, + }, + } + + const tool = createCompressMessageTool({ + client, + state, + logger, + config: buildConfig("message"), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressRange: "", compressMessage: "" } + }, + }, + } as any) + + await tool.execute( + { + topic: "Batch stale notes", + content: [ + { + messageId: "m0002", + topic: "Code path note", + summary: "Captured the assistant code-path findings.", + }, + { + messageId: "m0003", + topic: "Task output note", + summary: "Captured the assistant task-backed follow-up.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-message-group", + }, + ) + + appendOriginMessage(rawMessages, sessionID, "msg-compress-message-group") + + const blocks = Array.from(state.prune.messages.blocksById.values()).sort( + (a, b) => a.blockId - b.blockId, + ) + assert.equal(blocks.length, 2) + assert.equal(blocks[0]?.runId, blocks[1]?.runId) + assert.equal(blocks[0]?.batchTopic, "Batch stale notes") + + await handleDecompressCommand({ + client, + state, + logger, + sessionId: sessionID, + messages: rawMessages, + args: [], + }) + + const groupedListMessage = ignoredMessages.pop() || "" + assert.match(groupedListMessage, /Compression #1 - 2 messages - Batch stale notes/) + assert.doesNotMatch(groupedListMessage, /Code path note/) + + await handleDecompressCommand({ + client, + state, + logger, + sessionId: sessionID, + messages: rawMessages, + args: [String(blocks[0]?.blockId || 1)], + }) + + assert.ok(blocks.every((block) => block.deactivatedByUser)) + assert.ok(blocks.every((block) => !block.active)) + + await handleRecompressCommand({ + client, + state, + logger, + sessionId: sessionID, + messages: rawMessages, + args: [String(blocks[0]?.blockId || 1)], + }) + + assert.ok(blocks.every((block) => !block.deactivatedByUser)) + assert.ok(blocks.every((block) => block.active)) +}) + +test("decompress keeps batched ranges individually restorable", async () => { + const sessionID = `ses_range_individual_decompress_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const state = createSessionState() + const logger = new Logger(false) + const ignoredMessages: string[] = [] + const client = { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + prompt: async ({ body }: { body: { parts: Array<{ text: string }> } }) => { + ignoredMessages.push(body.parts[0]?.text || "") + }, + }, + } + + const tool = createCompressRangeTool({ + client, + state, + logger, + config: buildConfig("range"), + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressRange: "", compressMessage: "" } + }, + }, + } as any) + + await tool.execute( + { + topic: "Batch stale notes", + content: [ + { + startId: "m0001", + endId: "m0001", + summary: "Captured the opening user request.", + }, + { + startId: "m0002", + endId: "m0002", + summary: "Captured the assistant code-path findings.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-range-group", + }, + ) + + appendOriginMessage(rawMessages, sessionID, "msg-compress-range-group") + + const blocks = Array.from(state.prune.messages.blocksById.values()).sort( + (a, b) => a.blockId - b.blockId, + ) + assert.equal(blocks.length, 2) + assert.equal(blocks[0]?.runId, blocks[1]?.runId) + + await handleDecompressCommand({ + client, + state, + logger, + sessionId: sessionID, + messages: rawMessages, + args: [], + }) + + const listMessage = ignoredMessages.pop() || "" + assert.match(listMessage, /1 \(.+\)\s+Compression #1 - Batch stale notes/) + assert.match(listMessage, /2 \(.+\)\s+Compression #1 - Batch stale notes/) + + await handleDecompressCommand({ + client, + state, + logger, + sessionId: sessionID, + messages: rawMessages, + args: [String(blocks[0]?.blockId || 1)], + }) + + assert.equal(blocks[0]?.deactivatedByUser, true) + assert.equal(blocks[0]?.active, false) + assert.equal(blocks[1]?.active, true) + assert.equal(blocks[1]?.deactivatedByUser, false) +}) diff --git a/tests/message-priority.test.ts b/tests/message-priority.test.ts new file mode 100644 index 00000000..a05cd15f --- /dev/null +++ b/tests/message-priority.test.ts @@ -0,0 +1,463 @@ +import assert from "node:assert/strict" +import test from "node:test" +import type { PluginConfig } from "../lib/config" +import { createTextCompleteHandler } from "../lib/hooks" +import { Logger } from "../lib/logger" +import { assignMessageRefs } from "../lib/message-ids" +import { injectMessageIds } from "../lib/messages/inject/inject" +import { applyAnchoredNudges } from "../lib/messages/inject/utils" +import { prune } from "../lib/messages/prune" +import { buildPriorityMap } from "../lib/messages/priority" +import { stripHallucinationsFromString } from "../lib/messages/utils" +import { createSessionState, type WithParts } from "../lib/state" + +function buildConfig(mode: "message" | "range" = "message"): PluginConfig { + return { + enabled: true, + debug: false, + pruneNotification: "off", + pruneNotificationType: "chat", + commands: { + enabled: true, + protectedTools: [], + }, + manualMode: { + enabled: false, + automaticStrategies: true, + }, + turnProtection: { + enabled: false, + turns: 4, + }, + experimental: { + allowSubAgents: false, + customPrompts: false, + }, + protectedFilePatterns: [], + compress: { + mode, + permission: "allow", + showCompression: false, + maxContextLimit: 150000, + minContextLimit: 50000, + nudgeFrequency: 5, + iterationNudgeThreshold: 15, + nudgeForce: "soft", + protectedTools: ["task"], + protectUserMessages: false, + }, + strategies: { + deduplication: { + enabled: true, + protectedTools: [], + }, + purgeErrors: { + enabled: true, + turns: 4, + protectedTools: [], + }, + }, + } +} + +function textPart(messageID: string, sessionID: string, id: string, text: string) { + return { + id, + messageID, + sessionID, + type: "text" as const, + text, + } +} + +function toolPart( + messageID: string, + sessionID: string, + callID: string, + toolName: string, + output: string, +) { + return { + id: `${callID}-part`, + messageID, + sessionID, + type: "tool" as const, + tool: toolName, + callID, + state: { + status: "completed" as const, + input: { description: "demo" }, + output, + }, + } +} + +function buildMessage( + id: string, + role: "user" | "assistant", + sessionID: string, + text: string, + created: number, +): WithParts { + const info = + role === "user" + ? { + id, + role, + sessionID, + agent: "assistant", + model: { + providerID: "anthropic", + modelID: "claude-test", + }, + time: { created }, + } + : { + id, + role, + sessionID, + agent: "assistant", + time: { created }, + } + + return { + info: info as WithParts["info"], + parts: [textPart(id, sessionID, `${id}-part`, text)], + } +} + +function repeatedWord(word: string, count: number): string { + return Array.from({ length: count }, () => word).join(" ") +} + +test("injectMessageIds prefers assistant tool outputs over text parts in message mode", () => { + const sessionID = "ses_message_priority_tags" + const messages: WithParts[] = [ + buildMessage("msg-user-1", "user", sessionID, repeatedWord("investigate", 6000), 1), + { + info: { + id: "msg-assistant-1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 2 }, + } as WithParts["info"], + parts: [ + textPart( + "msg-assistant-1", + sessionID, + "msg-assistant-1-part", + "Short follow-up note.", + ), + toolPart("msg-assistant-1", sessionID, "call-task-1", "task", "task output body"), + ], + }, + ] + const state = createSessionState() + const config = buildConfig() + + assignMessageRefs(state, messages) + const compressionPriorities = buildPriorityMap(config, state, messages) + + injectMessageIds(state, config, messages, compressionPriorities) + + assert.equal(messages[0]?.parts.length, 1) + assert.equal(messages[1]?.parts.length, 2) + + const userText = messages[0]?.parts[0] + const assistantText = messages[1]?.parts[0] + const assistantTool = messages[1]?.parts[1] + + assert.equal(userText?.type, "text") + assert.equal(assistantText?.type, "text") + assert.equal(assistantTool?.type, "tool") + assert.match( + (userText as any).text, + /\n\nm0001<\/dcp-message-id>/, + ) + assert.equal((assistantText as any).text, "Short follow-up note.") + assert.match( + (assistantTool as any).state.output, + /m0002<\/dcp-message-id>/, + ) +}) + +test("injectMessageIds marks protected user messages as BLOCKED without priority in message mode", () => { + const sessionID = "ses_message_blocked_user_tags" + const messages: WithParts[] = [ + buildMessage("msg-user-1", "user", sessionID, repeatedWord("investigate", 6000), 1), + buildMessage("msg-assistant-1", "assistant", sessionID, "Short follow-up note.", 2), + ] + const state = createSessionState() + const config = buildConfig() + config.compress.protectUserMessages = true + + assignMessageRefs(state, messages) + const compressionPriorities = buildPriorityMap(config, state, messages) + + injectMessageIds(state, config, messages, compressionPriorities) + + const userText = messages[0]?.parts[0] + const assistantText = messages[1]?.parts[0] + + assert.equal(userText?.type, "text") + assert.equal(assistantText?.type, "text") + assert.match((userText as any).text, /\n\nBLOCKED<\/dcp-message-id>/) + assert.doesNotMatch((userText as any).text, /priority=/) + assert.match( + (assistantText as any).text, + /\n\nm0002<\/dcp-message-id>/, + ) +}) + +test("message-mode nudges append to existing text parts and list only earlier visible high-priority message IDs", () => { + const sessionID = "ses_message_priority_nudges" + const messages: WithParts[] = [ + buildMessage("msg-user-1", "user", sessionID, repeatedWord("alpha", 6000), 1), + buildMessage("msg-assistant-1", "assistant", sessionID, repeatedWord("beta", 6000), 2), + buildMessage("msg-user-2", "user", sessionID, repeatedWord("gamma", 6000), 3), + buildMessage("msg-assistant-2", "assistant", sessionID, repeatedWord("delta", 6000), 4), + ] + const state = createSessionState() + const config = buildConfig() + + assignMessageRefs(state, messages) + state.prune.messages.byMessageId.set("msg-assistant-1", { + tokenCount: 999, + allBlockIds: [1], + activeBlockIds: [1], + }) + state.nudges.contextLimitAnchors.add("msg-user-2") + + const compressionPriorities = buildPriorityMap(config, state, messages) + + applyAnchoredNudges( + state, + config, + messages, + { + system: "", + compressRange: "", + compressMessage: "", + contextLimitNudge: "Base context nudge", + turnNudge: "Base turn nudge", + iterationNudge: "Base iteration nudge", + }, + compressionPriorities, + ) + + assert.equal(messages[2]?.parts.length, 1) + + const injectedNudge = messages[2]?.parts[0] + assert.equal(injectedNudge?.type, "text") + assert.match((injectedNudge as any).text, /\n\nBase context nudge/) + assert.match((injectedNudge as any).text, /Message priority context:/) + assert.match((injectedNudge as any).text, /High-priority message IDs before this point: m0001/) + assert.doesNotMatch((injectedNudge as any).text, /m0002/) + assert.doesNotMatch((injectedNudge as any).text, /m0003/) + assert.doesNotMatch((injectedNudge as any).text, /m0004/) +}) + +test("message-mode nudges exclude protected user messages from priority guidance", () => { + const sessionID = "ses_message_blocked_priority_nudges" + const messages: WithParts[] = [ + buildMessage("msg-user-1", "user", sessionID, repeatedWord("alpha", 6000), 1), + buildMessage("msg-assistant-1", "assistant", sessionID, repeatedWord("beta", 6000), 2), + buildMessage("msg-user-2", "user", sessionID, repeatedWord("gamma", 6000), 3), + ] + const state = createSessionState() + const config = buildConfig() + config.compress.protectUserMessages = true + + assignMessageRefs(state, messages) + state.nudges.contextLimitAnchors.add("msg-user-2") + + const compressionPriorities = buildPriorityMap(config, state, messages) + + applyAnchoredNudges( + state, + config, + messages, + { + system: "", + compressRange: "", + compressMessage: "", + contextLimitNudge: "Base context nudge", + turnNudge: "Base turn nudge", + iterationNudge: "Base iteration nudge", + }, + compressionPriorities, + ) + + const injectedNudge = messages[2]?.parts[0] + assert.equal(injectedNudge?.type, "text") + assert.match((injectedNudge as any).text, /High-priority message IDs before this point: m0002/) + assert.doesNotMatch((injectedNudge as any).text, /m0001/) +}) + +test("range-mode nudges append to existing text parts before tool outputs", () => { + const sessionID = "ses_range_nudge_injection" + const messages: WithParts[] = [ + buildMessage("msg-user-1", "user", sessionID, repeatedWord("alpha", 6000), 1), + { + info: { + id: "msg-assistant-1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 2 }, + } as WithParts["info"], + parts: [ + textPart("msg-assistant-1", sessionID, "msg-assistant-1-part", "Working summary."), + toolPart("msg-assistant-1", sessionID, "call-task-2", "task", "task output body"), + ], + }, + ] + const state = createSessionState() + const config = buildConfig("range") + + assignMessageRefs(state, messages) + state.prune.messages.activeBlockIds.add(7) + state.nudges.contextLimitAnchors.add("msg-assistant-1") + + applyAnchoredNudges(state, config, messages, { + system: "", + compressRange: "", + compressMessage: "", + contextLimitNudge: "Base context nudge", + turnNudge: "Base turn nudge", + iterationNudge: "Base iteration nudge", + }) + + assert.equal(messages[1]?.parts.length, 2) + + const injectedNudge = messages[1]?.parts[0] + const toolOutput = messages[1]?.parts[1] + assert.equal(injectedNudge?.type, "text") + assert.equal(toolOutput?.type, "tool") + assert.match((injectedNudge as any).text, /\n\nBase context nudge/) + assert.match((injectedNudge as any).text, /Compressed block context:/) + assert.match((injectedNudge as any).text, /Active compressed blocks in this session: 1 \(b7\)/) + assert.equal((toolOutput as any).state.output, "task output body") +}) + +test("message-mode rendered compressed summaries mark block IDs as BLOCKED", () => { + const sessionID = "ses_message_blocked_blocks" + const messages: WithParts[] = [ + buildMessage("msg-user-1", "user", sessionID, "Original request", 1), + buildMessage("msg-assistant-1", "assistant", sessionID, "Follow-up", 2), + ] + const state = createSessionState() + const config = buildConfig("message") + const logger = new Logger(false) + + state.prune.messages.byMessageId.set("msg-user-1", { + tokenCount: 20, + allBlockIds: [7], + activeBlockIds: [7], + }) + state.prune.messages.blocksById.set(7, { + blockId: 7, + runId: 1, + active: true, + deactivatedByUser: false, + compressedTokens: 0, + mode: "range", + topic: "Earlier notes", + batchTopic: "Earlier notes", + startId: "m0001", + endId: "m0001", + anchorMessageId: "msg-user-1", + compressMessageId: "msg-origin", + includedBlockIds: [], + consumedBlockIds: [], + parentBlockIds: [], + directMessageIds: ["msg-user-1"], + directToolIds: [], + effectiveMessageIds: ["msg-user-1"], + effectiveToolIds: [], + createdAt: 1, + summary: + "[Compressed conversation section]\nEarlier summary\n\nb7", + }) + state.prune.messages.activeBlockIds.add(7) + state.prune.messages.activeByAnchorMessageId.set("msg-user-1", 7) + + prune(state, logger, config, messages) + + const summaryText = (messages[0]?.parts[0] as any)?.text || "" + assert.match(summaryText, /BLOCKED<\/dcp-message-id>/) + assert.doesNotMatch(summaryText, /b7<\/dcp-message-id>/) +}) + +test("range-mode rendered compressed summaries keep block IDs", () => { + const sessionID = "ses_range_visible_blocks" + const messages: WithParts[] = [ + buildMessage("msg-user-1", "user", sessionID, "Original request", 1), + buildMessage("msg-assistant-1", "assistant", sessionID, "Follow-up", 2), + ] + const state = createSessionState() + const config = buildConfig("range") + const logger = new Logger(false) + + state.prune.messages.byMessageId.set("msg-user-1", { + tokenCount: 20, + allBlockIds: [7], + activeBlockIds: [7], + }) + state.prune.messages.blocksById.set(7, { + blockId: 7, + runId: 1, + active: true, + deactivatedByUser: false, + compressedTokens: 0, + mode: "range", + topic: "Earlier notes", + batchTopic: "Earlier notes", + startId: "m0001", + endId: "m0001", + anchorMessageId: "msg-user-1", + compressMessageId: "msg-origin", + includedBlockIds: [], + consumedBlockIds: [], + parentBlockIds: [], + directMessageIds: ["msg-user-1"], + directToolIds: [], + effectiveMessageIds: ["msg-user-1"], + effectiveToolIds: [], + createdAt: 1, + summary: + "[Compressed conversation section]\nEarlier summary\n\nb7", + }) + state.prune.messages.activeBlockIds.add(7) + state.prune.messages.activeByAnchorMessageId.set("msg-user-1", 7) + + prune(state, logger, config, messages) + + const summaryText = (messages[0]?.parts[0] as any)?.text || "" + assert.match(summaryText, /b7<\/dcp-message-id>/) + assert.doesNotMatch(summaryText, /BLOCKED<\/dcp-message-id>/) +}) + +test("hallucination stripping removes exact metadata tags and preserves lookalikes", async () => { + const text = + 'alpham0007' + + "BLOCKED" + + 'm0008' + + 'remove this' + + "keep this" + + "omega" + + assert.equal( + stripHallucinationsFromString(text), + 'alpham0008keep thisomega', + ) + + const handler = createTextCompleteHandler() + const output = { text } + await handler({ sessionID: "session", messageID: "message", partID: "part" }, output) + assert.equal( + output.text, + 'alpham0008keep thisomega', + ) +}) diff --git a/tests/prompts.test.ts b/tests/prompts.test.ts index 9a3974c5..b1c0db0b 100644 --- a/tests/prompts.test.ts +++ b/tests/prompts.test.ts @@ -7,7 +7,7 @@ import { Logger } from "../lib/logger" import { PromptStore } from "../lib/prompts/store" import { SYSTEM as SYSTEM_PROMPT } from "../lib/prompts/system" -function createPromptStoreFixture(overrideContent?: string) { +function createPromptStoreFixture(overrideContent?: string, overrideFileName = "system.md") { const rootDir = mkdtempSync(join(tmpdir(), "opencode-dcp-prompts-")) const configHome = join(rootDir, "config") const workspaceDir = join(rootDir, "workspace") @@ -24,7 +24,7 @@ function createPromptStoreFixture(overrideContent?: string) { if (overrideContent !== undefined) { const overrideDir = join(configHome, "opencode", "dcp-prompts", "overrides") mkdirSync(overrideDir, { recursive: true }) - writeFileSync(join(overrideDir, "system.md"), overrideContent, "utf-8") + writeFileSync(join(overrideDir, overrideFileName), overrideContent, "utf-8") } const store = new PromptStore(new Logger(false), workspaceDir, true) @@ -101,3 +101,65 @@ test("system prompt overrides handle reminder tags safely", async (t) => { } }) }) + +test("prompt store exposes bundled message-mode compress prompt", () => { + const fixture = createPromptStoreFixture() + + try { + const runtimePrompts = fixture.store.getRuntimePrompts() + + assert.match(runtimePrompts.compressMessage, /selected individual messages/i) + assert.match( + runtimePrompts.compressMessage, + /Only use raw message IDs of the form `mNNNN`\./, + ) + assert.match(runtimePrompts.compressMessage, /priority="high"/) + assert.match(runtimePrompts.compressMessage, /prefer higher-priority messages first/i) + assert.match(runtimePrompts.compressMessage, /BLOCKED/) + assert.match(runtimePrompts.compressMessage, /cannot be compressed/i) + assert.match(runtimePrompts.compressMessage, /Do not use compressed block placeholders/i) + assert.doesNotMatch(runtimePrompts.compressMessage, /THE FORMAT OF COMPRESS/) + } finally { + fixture.cleanup() + } +}) + +test("compress-message overrides preserve plain-text metadata mentions", () => { + const fixture = createPromptStoreFixture( + [ + "Override body.", + "", + 'Each message has an ID inside XML metadata tags like `m0007`.', + "Messages marked as `BLOCKED` cannot be compressed.", + ].join("\n"), + "compress-message.md", + ) + + try { + const runtimePrompts = fixture.store.getRuntimePrompts() + + assert.match(runtimePrompts.compressMessage, /Override body\./) + assert.match( + runtimePrompts.compressMessage, + /m0007<\/dcp-message-id>/, + ) + assert.match(runtimePrompts.compressMessage, /BLOCKED<\/dcp-message-id>/) + } finally { + fixture.cleanup() + } +}) + +test("prompt store exposes bundled range-mode compress prompt", () => { + const fixture = createPromptStoreFixture() + + try { + const runtimePrompts = fixture.store.getRuntimePrompts() + + assert.match(runtimePrompts.compressRange, /Collapse a range in the conversation/i) + assert.match(runtimePrompts.compressRange, /COMPRESSED BLOCK PLACEHOLDERS/) + assert.match(runtimePrompts.compressRange, /BATCHING/) + assert.match(runtimePrompts.compressRange, /content` array/) + } finally { + fixture.cleanup() + } +})