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
54 changes: 54 additions & 0 deletions lib/compress/protected-content.ts
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,60 @@ 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
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") 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
16 changes: 15 additions & 1 deletion lib/update.ts
Original file line number Diff line number Diff line change
Expand Up @@ -88,12 +88,26 @@ export async function updateRemoveDir(packageDir: string, name: string) {

const wrapperDir = dirname(nodeModulesDir)
const wrapperPkg = await readPackageJson(join(wrapperDir, "package.json"))
const spec = wrapperPkg?.dependencies?.[name]
const spec = wrapperSpec(wrapperDir, name) ?? wrapperPkg?.dependencies?.[name]
if (!spec || !isAutoUpdatableSpec(spec)) return undefined

return wrapperDir
}

function wrapperSpec(wrapperDir: string, name: string) {
if (name.startsWith("@")) {
const [scope, pkg] = name.split("/")
if (!scope || !pkg || basename(dirname(wrapperDir)) !== scope) return undefined
const prefix = `${pkg}@`
const base = basename(wrapperDir)
return base.startsWith(prefix) ? base.slice(prefix.length) : undefined
}

const prefix = `${name}@`
const base = basename(wrapperDir)
return base.startsWith(prefix) ? base.slice(prefix.length) : undefined
}

export function isAutoUpdatableSpec(spec: string) {
const value = spec.trim()
if (!value) return false
Expand Down
118 changes: 118 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,123 @@ 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 user = rawMessages.find((message) => message.info.id === "msg-user-1")
const part = user?.parts[0]
if (part?.type === "text") {
part.text = "Investigate the issue. <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: "m0001",
topic: "User request note",
summary: "Captured the user's investigation request.",
},
],
},
{
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 ignores protect tags on ignored user messages", async () => {
const sessionID = `ses_message_ignored_protect_tag_${Date.now()}`
const rawMessages = buildMessages(sessionID)
const user = rawMessages.find((message) => message.info.id === "msg-user-1")
const part = user?.parts[0] as any
if (part?.type === "text") {
part.text = "Ignored notification. <protect>Do not preserve ignored note.</protect>"
part.ignored = true
}

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: "Ignored protected note",
content: [
{
messageId: "m0001",
topic: "Ignored note",
summary: "Captured the ignored user message.",
},
],
},
{
ask: async () => {},
metadata: () => {},
sessionID,
messageID: "msg-compress-ignored-protect-tag",
},
)

const block = Array.from(state.prune.messages.blocksById.values())[0]
assert.doesNotMatch(
block?.summary || "",
/The following protected prompt information was included in this conversation verbatim:/,
)
assert.doesNotMatch(block?.summary || "", /Do not preserve ignored note\./)
})

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
Loading
Loading