From f7aaeaef66eba551f58c058b89ec1477a7e94200 Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Thu, 7 May 2026 11:24:41 -0400 Subject: [PATCH] feat: add protect tags --- README.md | 2 + dcp.schema.json | 6 +++ lib/compress/message.ts | 12 ++++- lib/compress/protected-content.ts | 52 +++++++++++++++++++ lib/compress/range.ts | 16 +++++- lib/config.ts | 12 +++++ tests/compress-message.test.ts | 59 +++++++++++++++++++++ tests/compress-range.test.ts | 86 +++++++++++++++++++++++++++++++ tests/compression-groups.test.ts | 1 + tests/hooks-permission.test.ts | 1 + tests/message-priority.test.ts | 1 + tests/token-usage.test.ts | 1 + 12 files changed, 245 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 77902644..e995d65c 100644 --- a/README.md +++ b/README.md @@ -146,6 +146,8 @@ Each level overrides the previous, so project settings take priority over global "nudgeForce": "soft", // Tool names whose completed outputs are appended to the compression "protectedTools": [], + // Preserve text wrapped in ... when compressed + "protectTags": false, // Preserve your messages during compression. // Warning: large copy-pasted prompts will never be compressed away "protectUserMessages": false, diff --git a/dcp.schema.json b/dcp.schema.json index 6b37e4f6..39f2df53 100644 --- a/dcp.schema.json +++ b/dcp.schema.json @@ -237,6 +237,11 @@ "default": [], "description": "Tool names or wildcard patterns whose completed outputs should be appended to the compression summary. Supports glob wildcards: * matches any characters, ? matches a single character (e.g., \"mcp_*\", \"my_tool_?\")" }, + "protectTags": { + "type": "boolean", + "default": false, + "description": "Preserve text wrapped in ... when compressed" + }, "protectUserMessages": { "type": "boolean", "default": false, @@ -254,6 +259,7 @@ "iterationNudgeThreshold": 15, "nudgeForce": "soft", "protectedTools": [], + "protectTags": false, "protectUserMessages": false } }, diff --git a/lib/compress/message.ts b/lib/compress/message.ts index 5d81161e..d6bf8874 100644 --- a/lib/compress/message.ts +++ b/lib/compress/message.ts @@ -4,7 +4,7 @@ import { countTokens } from "../token-utils" import { MESSAGE_FORMAT_EXTENSION } from "../prompts/extensions/tool" import { formatIssues, formatResult, resolveMessages, validateArgs } from "./message-utils" import { finalizeSession, prepareSession, type NotificationEntry } from "./pipeline" -import { appendProtectedTools } from "./protected-content" +import { appendProtectedPromptInfo, appendProtectedTools } from "./protected-content" import { allocateBlockId, allocateRunId, @@ -77,11 +77,19 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType = [] for (const plan of plans) { + const summaryWithPromptInfo = appendProtectedPromptInfo( + plan.entry.summary, + plan.selection, + searchContext, + ctx.state, + ctx.config.compress.protectTags, + ) + const summaryWithTools = await appendProtectedTools( ctx.client, ctx.state, ctx.config.experimental.allowSubAgents, - plan.entry.summary, + summaryWithPromptInfo, plan.selection, searchContext, ctx.config.compress.protectedTools, diff --git a/lib/compress/protected-content.ts b/lib/compress/protected-content.ts index 5001ed40..1af9ba5a 100644 --- a/lib/compress/protected-content.ts +++ b/lib/compress/protected-content.ts @@ -53,6 +53,58 @@ export function appendProtectedUserMessages( return summary + heading + body } +export function appendProtectedPromptInfo( + summary: string, + selection: SelectionResolution, + searchContext: SearchContext, + state: SessionState, + enabled: boolean, +): string { + if (!enabled) return summary + + const protectedTexts: 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 !== "text" || typeof part.text !== "string") continue + + protectedTexts.push(...extractProtectedPromptInfo(part.text)) + } + } + + if (protectedTexts.length === 0) { + return summary + } + + const heading = + "\n\nThe following protected prompt information was included in this conversation verbatim:" + const body = protectedTexts.map((text) => `\n${text}`).join("") + return summary + heading + body +} + +export function extractProtectedPromptInfo(text: string): string[] { + const protectedTexts: string[] = [] + const protectTagRegex = /([\s\S]*?)<\/protect>/gi + + for (const match of text.matchAll(protectTagRegex)) { + const protectedText = match[1]?.trim() + if (protectedText) { + protectedTexts.push(protectedText) + } + } + + return protectedTexts +} + export async function appendProtectedTools( client: any, state: SessionState, diff --git a/lib/compress/range.ts b/lib/compress/range.ts index cc2ceaa9..d320be89 100644 --- a/lib/compress/range.ts +++ b/lib/compress/range.ts @@ -3,7 +3,11 @@ import type { ToolContext } from "./types" import { countTokens } from "../token-utils" import { RANGE_FORMAT_EXTENSION } from "../prompts/extensions/tool" import { finalizeSession, prepareSession, type NotificationEntry } from "./pipeline" -import { appendProtectedTools, appendProtectedUserMessages } from "./protected-content" +import { + appendProtectedPromptInfo, + appendProtectedTools, + appendProtectedUserMessages, +} from "./protected-content" import { appendMissingBlockSummaries, injectBlockPlaceholders, @@ -108,11 +112,19 @@ export function createCompressRangeTool(ctx: ToolContext): ReturnType): ValidationErro }) } + if (compress.protectTags !== undefined && typeof compress.protectTags !== "boolean") { + errors.push({ + key: "compress.protectTags", + expected: "boolean", + actual: typeof compress.protectTags, + }) + } + if ( compress.protectUserMessages !== undefined && typeof compress.protectUserMessages !== "boolean" @@ -677,6 +687,7 @@ const defaultConfig: PluginConfig = { iterationNudgeThreshold: 15, nudgeForce: "soft", protectedTools: [...COMPRESS_DEFAULT_PROTECTED_TOOLS], + protectTags: false, protectUserMessages: false, }, strategies: { @@ -842,6 +853,7 @@ function mergeCompress( iterationNudgeThreshold: override.iterationNudgeThreshold ?? base.iterationNudgeThreshold, nudgeForce: override.nudgeForce ?? base.nudgeForce, protectedTools: [...new Set([...base.protectedTools, ...(override.protectedTools ?? [])])], + protectTags: override.protectTags ?? base.protectTags, protectUserMessages: override.protectUserMessages ?? base.protectUserMessages, } } diff --git a/tests/compress-message.test.ts b/tests/compress-message.test.ts index ad0ab394..b1d9dd53 100644 --- a/tests/compress-message.test.ts +++ b/tests/compress-message.test.ts @@ -50,6 +50,7 @@ function buildConfig(): PluginConfig { iterationNudgeThreshold: 15, nudgeForce: "soft", protectedTools: ["task"], + protectTags: false, protectUserMessages: false, }, strategies: { @@ -226,6 +227,64 @@ test("compress message mode batches individual message summaries", async () => { assert.match(blocks[1]?.summary || "", /task output body/) }) +test("compress message mode appends protected prompt info", async () => { + const sessionID = `ses_message_protect_tag_${Date.now()}` + const rawMessages = buildMessages(sessionID) + const assistant = rawMessages.find((message) => message.info.id === "msg-assistant-1") + const part = assistant?.parts[0] + if (part?.type === "text") { + part.text = "I mapped the code path. Always preserve release checklist." + } + + const state = createSessionState() + const logger = new Logger(false) + const config = buildConfig() + config.compress.protectTags = 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) + + await tool.execute( + { + topic: "Protected note", + content: [ + { + messageId: "m0002", + topic: "Code path note", + summary: "Captured the assistant's code-path findings.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-protect-tag", + }, + ) + + const block = Array.from(state.prune.messages.blocksById.values())[0] + assert.match( + block?.summary || "", + /The following protected prompt information was included in this conversation verbatim:/, + ) + assert.match(block?.summary || "", /Always preserve release checklist\./) +}) + test("compress message mode stores call id for later duration attachment", async () => { const sessionID = `ses_message_compress_duration_${Date.now()}` const rawMessages = buildMessages(sessionID) diff --git a/tests/compress-range.test.ts b/tests/compress-range.test.ts index 7899189e..ff9c7161 100644 --- a/tests/compress-range.test.ts +++ b/tests/compress-range.test.ts @@ -50,6 +50,7 @@ function buildConfig(): PluginConfig { iterationNudgeThreshold: 15, nudgeForce: "soft", protectedTools: [], + protectTags: false, protectUserMessages: false, }, strategies: { @@ -178,6 +179,91 @@ test("compress range rebuilds subagent message refs after session state was rese assert.equal(state.prune.messages.blocksById.size, 1) }) +test("compress range mode appends protected prompt info", async () => { + const sessionID = `ses_range_protect_tag_${Date.now()}` + const rawMessages: WithParts[] = [ + { + 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-user-1", + "Investigate the release. Keep the npm publish token note.", + ), + ], + }, + { + info: { + id: "msg-assistant-1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 2 }, + } as WithParts["info"], + parts: [textPart("msg-assistant-1", sessionID, "part-assistant-1", "I checked it")], + }, + ] + + const state = createSessionState() + const logger = new Logger(false) + const config = buildConfig() + config.compress.protectTags = true + const tool = createCompressRangeTool({ + client: { + session: { + messages: async () => ({ data: rawMessages }), + get: async () => ({ data: { parentID: null } }), + }, + }, + state, + logger, + config, + prompts: { + reload() {}, + getRuntimePrompts() { + return { compressRange: "", compressMessage: "" } + }, + }, + } as any) + + await tool.execute( + { + topic: "Protected range", + content: [ + { + startId: "m0001", + endId: "m0002", + summary: "Captured release investigation.", + }, + ], + }, + { + ask: async () => {}, + metadata: () => {}, + sessionID, + messageID: "msg-compress-range-protect-tag", + }, + ) + + const block = Array.from(state.prune.messages.blocksById.values())[0] + assert.match( + block?.summary || "", + /The following protected prompt information was included in this conversation verbatim:/, + ) + assert.match(block?.summary || "", /Keep the npm publish token note\./) +}) + test("compress range mode batches multiple ranges into one notification", async () => { const sessionID = `ses_range_compress_batch_${Date.now()}` const rawMessages = buildMessages(sessionID) diff --git a/tests/compression-groups.test.ts b/tests/compression-groups.test.ts index f97d02ba..c5bac056 100644 --- a/tests/compression-groups.test.ts +++ b/tests/compression-groups.test.ts @@ -53,6 +53,7 @@ function buildConfig(mode: "message" | "range"): PluginConfig { iterationNudgeThreshold: 15, nudgeForce: "soft", protectedTools: ["task"], + protectTags: false, protectUserMessages: false, }, strategies: { diff --git a/tests/hooks-permission.test.ts b/tests/hooks-permission.test.ts index 5cfd5bf1..8364a6fc 100644 --- a/tests/hooks-permission.test.ts +++ b/tests/hooks-permission.test.ts @@ -48,6 +48,7 @@ function buildConfig(permission: "allow" | "ask" | "deny" = "allow"): PluginConf iterationNudgeThreshold: 15, nudgeForce: "soft", protectedTools: ["task"], + protectTags: false, protectUserMessages: false, }, strategies: { diff --git a/tests/message-priority.test.ts b/tests/message-priority.test.ts index 748c2c2d..1342ca1b 100644 --- a/tests/message-priority.test.ts +++ b/tests/message-priority.test.ts @@ -44,6 +44,7 @@ function buildConfig(mode: "message" | "range" = "message"): PluginConfig { iterationNudgeThreshold: 15, nudgeForce: "soft", protectedTools: ["task"], + protectTags: false, protectUserMessages: false, }, strategies: { diff --git a/tests/token-usage.test.ts b/tests/token-usage.test.ts index 4c1e6af7..549edeae 100644 --- a/tests/token-usage.test.ts +++ b/tests/token-usage.test.ts @@ -41,6 +41,7 @@ function buildConfig(maxContextLimit: number, minContextLimit = 1): PluginConfig iterationNudgeThreshold: 15, nudgeForce: "soft", protectedTools: ["task"], + protectTags: false, protectUserMessages: false, }, strategies: {