Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 <protect>...</protect> when compressed
"protectTags": false,
// Preserve your messages during compression.
// Warning: large copy-pasted prompts will never be compressed away
"protectUserMessages": false,
Expand Down
6 changes: 6 additions & 0 deletions dcp.schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -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 <protect>...</protect> when compressed"
},
"protectUserMessages": {
"type": "boolean",
"default": false,
Expand All @@ -254,6 +259,7 @@
"iterationNudgeThreshold": 15,
"nudgeForce": "soft",
"protectedTools": [],
"protectTags": false,
"protectUserMessages": false
}
},
Expand Down
12 changes: 10 additions & 2 deletions lib/compress/message.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -77,11 +77,19 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
}> = []

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,
Expand Down
52 changes: 52 additions & 0 deletions lib/compress/protected-content.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = /<protect>([\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,
Expand Down
16 changes: 14 additions & 2 deletions lib/compress/range.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -108,11 +112,19 @@ export function createCompressRangeTool(ctx: ToolContext): ReturnType<typeof too
ctx.config.compress.protectUserMessages,
)

const summaryWithPromptInfo = appendProtectedPromptInfo(
summaryWithUsers,
plan.selection,
searchContext,
ctx.state,
ctx.config.compress.protectTags,
)

const summaryWithTools = await appendProtectedTools(
ctx.client,
ctx.state,
ctx.config.experimental.allowSubAgents,
summaryWithUsers,
summaryWithPromptInfo,
plan.selection,
searchContext,
ctx.config.compress.protectedTools,
Expand Down
12 changes: 12 additions & 0 deletions lib/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ export interface CompressConfig {
iterationNudgeThreshold: number
nudgeForce: "strong" | "soft"
protectedTools: string[]
protectTags: boolean
protectUserMessages: boolean
}

Expand Down Expand Up @@ -123,6 +124,7 @@ export const VALID_CONFIG_KEYS = new Set([
"compress.iterationNudgeThreshold",
"compress.nudgeForce",
"compress.protectedTools",
"compress.protectTags",
"compress.protectUserMessages",
"strategies",
"strategies.deduplication",
Expand Down Expand Up @@ -422,6 +424,14 @@ export function validateConfigTypes(config: Record<string, any>): 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"
Expand Down Expand Up @@ -677,6 +687,7 @@ const defaultConfig: PluginConfig = {
iterationNudgeThreshold: 15,
nudgeForce: "soft",
protectedTools: [...COMPRESS_DEFAULT_PROTECTED_TOOLS],
protectTags: false,
protectUserMessages: false,
},
strategies: {
Expand Down Expand Up @@ -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,
}
}
Expand Down
59 changes: 59 additions & 0 deletions tests/compress-message.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ function buildConfig(): PluginConfig {
iterationNudgeThreshold: 15,
nudgeForce: "soft",
protectedTools: ["task"],
protectTags: false,
protectUserMessages: false,
},
strategies: {
Expand Down Expand Up @@ -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. <protect>Always preserve release checklist.</protect>"
}

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)
Expand Down
86 changes: 86 additions & 0 deletions tests/compress-range.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ function buildConfig(): PluginConfig {
iterationNudgeThreshold: 15,
nudgeForce: "soft",
protectedTools: [],
protectTags: false,
protectUserMessages: false,
},
strategies: {
Expand Down Expand Up @@ -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. <protect>Keep the npm publish token note.</protect>",
),
],
},
{
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)
Expand Down
1 change: 1 addition & 0 deletions tests/compression-groups.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ function buildConfig(mode: "message" | "range"): PluginConfig {
iterationNudgeThreshold: 15,
nudgeForce: "soft",
protectedTools: ["task"],
protectTags: false,
protectUserMessages: false,
},
strategies: {
Expand Down
1 change: 1 addition & 0 deletions tests/hooks-permission.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ function buildConfig(permission: "allow" | "ask" | "deny" = "allow"): PluginConf
iterationNudgeThreshold: 15,
nudgeForce: "soft",
protectedTools: ["task"],
protectTags: false,
protectUserMessages: false,
},
strategies: {
Expand Down
1 change: 1 addition & 0 deletions tests/message-priority.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ function buildConfig(mode: "message" | "range" = "message"): PluginConfig {
iterationNudgeThreshold: 15,
nudgeForce: "soft",
protectedTools: ["task"],
protectTags: false,
protectUserMessages: false,
},
strategies: {
Expand Down
Loading
Loading