Skip to content

Commit 41d1aea

Browse files
TarquinenTarquinen
committed
prioritize high-cost message compression
1 parent 9c8aa18 commit 41d1aea

11 files changed

Lines changed: 513 additions & 59 deletions

‎lib/hooks.ts‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type { SessionState, WithParts } from "./state"
22
import type { Logger } from "./logger"
33
import type { PluginConfig } from "./config"
44
import { assignMessageRefs } from "./message-ids"
5+
import { buildPriorityMap } from "./messages/priority"
56
import { syncToolCache } from "./state/tool-cache"
67
import {
78
prune,
@@ -11,7 +12,12 @@ import {
1112
injectExtendedSubAgentResults,
1213
stripStaleMetadata,
1314
} from "./messages"
14-
import { buildToolIdList, isIgnoredUserMessage, stripHallucinations } from "./messages/utils"
15+
import {
16+
buildToolIdList,
17+
isIgnoredUserMessage,
18+
stripHallucinations,
19+
stripHallucinationsFromString,
20+
} from "./messages/utils"
1521
import { checkSession } from "./state"
1622
import { renderSystemPrompt } from "./prompts"
1723
import { handleStatsCommand } from "./commands/stats"
@@ -33,9 +39,6 @@ const INTERNAL_AGENT_SIGNATURES = [
3339
"Summarize what was done in this conversation",
3440
]
3541

36-
const DCP_MESSAGE_ID_TAG_REGEX = /<dcp-message-id>(?:m\d+|b\d+)<\/dcp-message-id>/g
37-
const DCP_SYSTEM_REMINDER_REGEX = /<dcp-system-reminder\b[^>]*>[\s\S]*?<\/dcp-system-reminder>/g
38-
3942
function applyManualPrompt(state: SessionState, messages: WithParts[], logger: Logger): void {
4043
const pending = state.pendingManualTrigger
4144
if (!pending) {
@@ -148,9 +151,17 @@ export function createChatMessageTransformHandler(
148151
output.messages,
149152
config.experimental.allowSubAgents,
150153
)
154+
const compressionPriorities = buildPriorityMap(config, state, output.messages)
151155
prompts.reload()
152-
injectCompressNudges(state, config, logger, output.messages, prompts.getRuntimePrompts())
153-
injectMessageIds(state, config, output.messages)
156+
injectCompressNudges(
157+
state,
158+
config,
159+
logger,
160+
output.messages,
161+
prompts.getRuntimePrompts(),
162+
compressionPriorities,
163+
)
164+
injectMessageIds(state, config, output.messages, compressionPriorities)
154165
applyManualPrompt(state, output.messages, logger)
155166
stripStaleMetadata(output.messages)
156167

@@ -280,8 +291,6 @@ export function createTextCompleteHandler() {
280291
_input: { sessionID: string; messageID: string; partID: string },
281292
output: { text: string },
282293
) => {
283-
output.text = output.text
284-
.replace(DCP_SYSTEM_REMINDER_REGEX, "")
285-
.replace(DCP_MESSAGE_ID_TAG_REGEX, "")
294+
output.text = stripHallucinationsFromString(output.text)
286295
}
287296
}

‎lib/message-ids.ts‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,30 @@ export function parseBoundaryId(id: string): ParsedBoundaryId | null {
9090
return null
9191
}
9292

93-
export function formatMessageIdTag(ref: string): string {
94-
return `\n<${MESSAGE_ID_TAG_NAME}>${ref}</${MESSAGE_ID_TAG_NAME}>`
93+
function escapeXmlAttribute(value: string): string {
94+
return value
95+
.replace(/&/g, "&amp;")
96+
.replace(/"/g, "&quot;")
97+
.replace(/</g, "&lt;")
98+
.replace(/>/g, "&gt;")
99+
}
100+
101+
export function formatMessageIdTag(
102+
ref: string,
103+
attributes?: Record<string, string | undefined>,
104+
): string {
105+
const serializedAttributes = Object.entries(attributes || {})
106+
.sort(([left], [right]) => left.localeCompare(right))
107+
.map(([name, value]) => {
108+
if (name.trim().length === 0 || typeof value !== "string" || value.length === 0) {
109+
return ""
110+
}
111+
112+
return ` ${name}="${escapeXmlAttribute(value)}"`
113+
})
114+
.join("")
115+
116+
return `\n<${MESSAGE_ID_TAG_NAME}${serializedAttributes}>${ref}</${MESSAGE_ID_TAG_NAME}>`
95117
}
96118

97119
export function assignMessageRefs(state: SessionState, messages: WithParts[]): number {

‎lib/messages/inject/inject.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import type { Logger } from "../../logger"
33
import type { PluginConfig } from "../../config"
44
import type { RuntimePrompts } from "../../prompts/store"
55
import { formatMessageIdTag } from "../../message-ids"
6+
import type { CompressionPriorityMap } from "../priority"
67
import { compressPermission, getLastUserMessage } from "../../shared-utils"
78
import { saveSessionState } from "../../state/persistence"
89
import {
@@ -29,6 +30,7 @@ export const injectCompressNudges = (
2930
logger: Logger,
3031
messages: WithParts[],
3132
prompts: RuntimePrompts,
33+
compressionPriorities?: CompressionPriorityMap,
3234
): void => {
3335
if (compressPermission(state, config) === "deny") {
3436
return
@@ -127,7 +129,7 @@ export const injectCompressNudges = (
127129
}
128130
}
129131

130-
applyAnchoredNudges(state, config, messages, prompts)
132+
applyAnchoredNudges(state, config, messages, prompts, compressionPriorities)
131133

132134
if (anchorsChanged) {
133135
void saveSessionState(state, logger)
@@ -138,6 +140,7 @@ export const injectMessageIds = (
138140
state: SessionState,
139141
config: PluginConfig,
140142
messages: WithParts[],
143+
compressionPriorities?: CompressionPriorityMap,
141144
): void => {
142145
if (compressPermission(state, config) === "deny") {
143146
return
@@ -153,7 +156,11 @@ export const injectMessageIds = (
153156
continue
154157
}
155158

156-
const tag = formatMessageIdTag(messageRef)
159+
const priority =
160+
config.compress.mode === "message"
161+
? compressionPriorities?.get(message.info.id)?.priority
162+
: undefined
163+
const tag = formatMessageIdTag(messageRef, priority ? { priority } : undefined)
157164

158165
if (message.info.role === "user") {
159166
message.parts.push(createSyntheticTextPart(message, tag))

‎lib/messages/inject/utils.ts‎

Lines changed: 144 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,19 @@
11
import type { SessionState, WithParts } from "../../state"
22
import type { PluginConfig } from "../../config"
3+
import { renderMessagePriorityGuidance } from "../../prompts/message-priority-guidance"
34
import type { RuntimePrompts } from "../../prompts/store"
45
import type { UserMessage } from "@opencode-ai/sdk/v2"
6+
import {
7+
type CompressionPriorityMap,
8+
type MessagePriority,
9+
listPriorityRefsBeforeIndex,
10+
} from "../priority"
511
import { createSyntheticTextPart, isIgnoredUserMessage } from "../utils"
612
import { getLastUserMessage } from "../../shared-utils"
713
import { getCurrentTokenUsage } from "../../strategies/utils"
814

15+
const MESSAGE_MODE_NUDGE_PRIORITY: MessagePriority = "high"
16+
917
export interface LastUserModelContext {
1018
providerId: string | undefined
1119
modelId: string | undefined
@@ -201,60 +209,73 @@ function appendGuidanceToDcpTag(hintText: string, guidance: string): string {
201209
return `${beforeClose}\n\n${guidance}\n${afterClose}`
202210
}
203211

204-
function applyAnchoredNudge(
205-
anchorMessageIds: Set<string>,
212+
function buildMessagePriorityGuidance(
206213
messages: WithParts[],
207-
hintText: string,
208-
): void {
209-
if (anchorMessageIds.size === 0) {
214+
compressionPriorities: CompressionPriorityMap | undefined,
215+
anchorIndex: number,
216+
priority: MessagePriority,
217+
): string {
218+
if (!compressionPriorities || compressionPriorities.size === 0) {
219+
return ""
220+
}
221+
222+
const refs = listPriorityRefsBeforeIndex(messages, compressionPriorities, anchorIndex, priority)
223+
const priorityLabel = `${priority[0].toUpperCase()}${priority.slice(1)}`
224+
225+
return renderMessagePriorityGuidance(priorityLabel, refs)
226+
}
227+
228+
function injectAnchoredNudge(message: WithParts, hintText: string): void {
229+
if (!hintText.trim()) {
210230
return
211231
}
212232

213-
for (const anchorMessageId of anchorMessageIds) {
214-
const messageIndex = messages.findIndex((message) => message.info.id === anchorMessageId)
215-
if (messageIndex === -1) {
216-
continue
217-
}
233+
if (message.info.role === "user") {
234+
message.parts.push(createSyntheticTextPart(message, hintText))
235+
return
236+
}
218237

219-
const message = messages[messageIndex]
220-
if (message.info.role === "user") {
221-
message.parts.push(createSyntheticTextPart(message, hintText))
222-
continue
223-
}
238+
if (message.info.role !== "assistant") {
239+
return
240+
}
224241

225-
if (message.info.role !== "assistant") {
242+
const syntheticPart = createSyntheticTextPart(message, hintText)
243+
const firstToolIndex = message.parts.findIndex((p) => p.type === "tool")
244+
if (firstToolIndex === -1) {
245+
message.parts.push(syntheticPart)
246+
} else {
247+
message.parts.splice(firstToolIndex, 0, syntheticPart)
248+
}
249+
}
250+
251+
function collectAnchoredMessages(
252+
anchorMessageIds: Set<string>,
253+
messages: WithParts[],
254+
): Array<{ message: WithParts; index: number }> {
255+
const anchoredMessages: Array<{ message: WithParts; index: number }> = []
256+
257+
for (const anchorMessageId of anchorMessageIds) {
258+
const index = messages.findIndex((message) => message.info.id === anchorMessageId)
259+
if (index === -1) {
226260
continue
227261
}
228262

229-
const syntheticPart = createSyntheticTextPart(message, hintText)
230-
const firstToolIndex = message.parts.findIndex((p) => p.type === "tool")
231-
if (firstToolIndex === -1) {
232-
message.parts.push(syntheticPart)
233-
} else {
234-
message.parts.splice(firstToolIndex, 0, syntheticPart)
235-
}
263+
anchoredMessages.push({
264+
message: messages[index],
265+
index,
266+
})
236267
}
268+
269+
return anchoredMessages
237270
}
238271

239-
export function applyAnchoredNudges(
272+
function collectTurnNudgeAnchors(
240273
state: SessionState,
241274
config: PluginConfig,
242275
messages: WithParts[],
243-
prompts: RuntimePrompts,
244-
): void {
245-
const compressedBlockGuidance =
246-
config.compress.mode === "message" ? "" : buildCompressedBlockGuidance(state)
247-
248-
const contextLimitNudge = appendGuidanceToDcpTag(
249-
prompts.contextLimitNudge,
250-
compressedBlockGuidance,
251-
)
252-
253-
applyAnchoredNudge(state.nudges.contextLimitAnchors, messages, contextLimitNudge)
254-
276+
): Set<string> {
255277
const turnNudgeAnchors = new Set<string>()
256278
const targetRole = config.compress.nudgeForce === "strong" ? "user" : "assistant"
257-
const turnNudge = appendGuidanceToDcpTag(prompts.turnNudge, compressedBlockGuidance)
258279

259280
for (const message of messages) {
260281
if (!state.nudges.turnNudgeAnchors.has(message.info.id)) continue
@@ -264,8 +285,91 @@ export function applyAnchoredNudges(
264285
}
265286
}
266287

267-
applyAnchoredNudge(turnNudgeAnchors, messages, turnNudge)
288+
return turnNudgeAnchors
289+
}
290+
291+
function applyRangeModeAnchoredNudge(
292+
anchorMessageIds: Set<string>,
293+
messages: WithParts[],
294+
basePrompt: string,
295+
compressedBlockGuidance: string,
296+
): void {
297+
const hintText = appendGuidanceToDcpTag(basePrompt, compressedBlockGuidance)
298+
if (!hintText.trim()) {
299+
return
300+
}
301+
302+
for (const { message } of collectAnchoredMessages(anchorMessageIds, messages)) {
303+
injectAnchoredNudge(message, hintText)
304+
}
305+
}
306+
307+
function applyMessageModeAnchoredNudge(
308+
anchorMessageIds: Set<string>,
309+
messages: WithParts[],
310+
basePrompt: string,
311+
compressionPriorities?: CompressionPriorityMap,
312+
): void {
313+
for (const { message, index } of collectAnchoredMessages(anchorMessageIds, messages)) {
314+
const priorityGuidance = buildMessagePriorityGuidance(
315+
messages,
316+
compressionPriorities,
317+
index,
318+
MESSAGE_MODE_NUDGE_PRIORITY,
319+
)
320+
const hintText = appendGuidanceToDcpTag(basePrompt, priorityGuidance)
321+
injectAnchoredNudge(message, hintText)
322+
}
323+
}
268324

269-
const iterationNudge = appendGuidanceToDcpTag(prompts.iterationNudge, compressedBlockGuidance)
270-
applyAnchoredNudge(state.nudges.iterationNudgeAnchors, messages, iterationNudge)
325+
export function applyAnchoredNudges(
326+
state: SessionState,
327+
config: PluginConfig,
328+
messages: WithParts[],
329+
prompts: RuntimePrompts,
330+
compressionPriorities?: CompressionPriorityMap,
331+
): void {
332+
const turnNudgeAnchors = collectTurnNudgeAnchors(state, config, messages)
333+
334+
if (config.compress.mode === "message") {
335+
applyMessageModeAnchoredNudge(
336+
state.nudges.contextLimitAnchors,
337+
messages,
338+
prompts.contextLimitNudge,
339+
compressionPriorities,
340+
)
341+
applyMessageModeAnchoredNudge(
342+
turnNudgeAnchors,
343+
messages,
344+
prompts.turnNudge,
345+
compressionPriorities,
346+
)
347+
applyMessageModeAnchoredNudge(
348+
state.nudges.iterationNudgeAnchors,
349+
messages,
350+
prompts.iterationNudge,
351+
compressionPriorities,
352+
)
353+
return
354+
}
355+
356+
const compressedBlockGuidance = buildCompressedBlockGuidance(state)
357+
applyRangeModeAnchoredNudge(
358+
state.nudges.contextLimitAnchors,
359+
messages,
360+
prompts.contextLimitNudge,
361+
compressedBlockGuidance,
362+
)
363+
applyRangeModeAnchoredNudge(
364+
turnNudgeAnchors,
365+
messages,
366+
prompts.turnNudge,
367+
compressedBlockGuidance,
368+
)
369+
applyRangeModeAnchoredNudge(
370+
state.nudges.iterationNudgeAnchors,
371+
messages,
372+
prompts.iterationNudge,
373+
compressedBlockGuidance,
374+
)
271375
}

0 commit comments

Comments
 (0)