prioritize compress tool messages

Keep completed compress calls visible as high-priority message targets so message-mode guidance preserves them consistently.

Author: Tarquinen

lib/messages/inject/inject.ts

@@ -4,7 +4,7 @@ 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 { compressPermission, getLastUserMessage, messageHasCompress } from "../../shared-utils"
 import { saveSessionState } from "../../state/persistence"
 import {
     appendToTextPart,
@@ -22,7 +22,6 @@ import {
     getNudgeFrequency,
     getModelInfo,
     isContextOverLimits,
-    messageHasCompress,
 } from "./utils"
 
 export const injectCompressNudges = (

lib/messages/inject/utils.ts

@@ -58,14 +58,6 @@ export function countMessagesAfterIndex(messages: WithParts[], index: number): n
     return count
 }
 
-export function messageHasCompress(message: WithParts): boolean {
-    const parts = Array.isArray(message.parts) ? message.parts : []
-    return parts.some(
-        (part) =>
-            part.type === "tool" && part.state.status === "completed" && part.tool === "compress",
-    )
-}
-
 export function getModelInfo(messages: WithParts[]): LastUserModelContext {
     const lastUserMessage = getLastUserMessage(messages)
     if (!lastUserMessage) {

lib/messages/priority.ts

@@ -1,6 +1,6 @@
 import type { PluginConfig } from "../config"
 import { countAllMessageTokens } from "../strategies/utils"
-import { isMessageCompacted } from "../shared-utils"
+import { isMessageCompacted, messageHasCompress } from "../shared-utils"
 import type { SessionState, WithParts } from "../state"
 import { isIgnoredUserMessage, isProtectedUserMessage } from "./utils"
 
@@ -54,7 +54,7 @@ export function buildPriorityMap(
         priorities.set(rawMessageId, {
             ref,
             tokenCount,
-            priority: classifyMessagePriority(tokenCount),
+            priority: messageHasCompress(message) ? "high" : classifyMessagePriority(tokenCount),
         })
     }
 

lib/shared-utils.ts

@@ -28,6 +28,18 @@ export const getLastUserMessage = (
     return null
 }
 
+export const messageHasCompress = (message: WithParts): boolean => {
+    if (message.info.role !== "assistant") {
+        return false
+    }
+
+    const parts = Array.isArray(message.parts) ? message.parts : []
+    return parts.some(
+        (part) =>
+            part.type === "tool" && part.tool === "compress" && part.state?.status === "completed",
+    )
+}
+
 export const compressPermission = (
     state: SessionState,
     config: PluginConfig,

lib/state/utils.ts

@@ -5,7 +5,7 @@ import type {
     SessionState,
     WithParts,
 } from "./types"
-import { isMessageCompacted } from "../shared-utils"
+import { isMessageCompacted, messageHasCompress } from "../shared-utils"
 import { isIgnoredUserMessage } from "../messages/utils"
 
 interface PersistedPruneMessagesState {
@@ -236,26 +236,14 @@ export function loadPruneMessagesState(
     return state
 }
 
-function hasCompletedCompress(message: WithParts): boolean {
-    if (message.info.role !== "assistant") {
-        return false
-    }
-
-    const parts = Array.isArray(message.parts) ? message.parts : []
-    return parts.some(
-        (part) =>
-            part.type === "tool" && part.tool === "compress" && part.state?.status === "completed",
-    )
-}
-
 export function collectTurnNudgeAnchors(messages: WithParts[]): Set<string> {
     const anchors = new Set<string>()
     let pendingUserMessageId: string | null = null
 
     for (let i = messages.length - 1; i >= 0; i--) {
         const message = messages[i]
 
-        if (hasCompletedCompress(message)) {
+        if (messageHasCompress(message)) {
             break
         }
 

tests/message-priority.test.ts

@@ -333,6 +333,53 @@ test("injectMessageIds tags every text part and tool output in range mode", () =
     assert.match((assistantToolTwo as any).state.output, /<dcp-message-id>m0002<\/dcp-message-id>/)
 })
 
+test("message mode marks compress tool messages as high priority even when short", () => {
+    const sessionID = "ses_message_compress_high_priority"
+    const messages: WithParts[] = [
+        buildMessage("msg-user-1", "user", sessionID, "Please compress this chunk.", 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-1", "Done."),
+                toolPart(
+                    "msg-assistant-1",
+                    sessionID,
+                    "call-compress-1",
+                    "compress",
+                    "[Compressed conversation section]",
+                ),
+            ],
+        },
+    ]
+    const state = createSessionState()
+    const config = buildConfig()
+
+    assignMessageRefs(state, messages)
+    const compressionPriorities = buildPriorityMap(config, state, messages)
+
+    assert.equal(compressionPriorities.get("msg-assistant-1")?.priority, "high")
+
+    injectMessageIds(state, config, messages, compressionPriorities)
+
+    const assistantText = messages[1]?.parts[0]
+    const assistantTool = messages[1]?.parts[1]
+
+    assert.match(
+        (assistantText as any).text,
+        /\n\n<dcp-message-id priority="high">m0002<\/dcp-message-id>/,
+    )
+    assert.match(
+        (assistantTool as any).state.output,
+        /<dcp-message-id priority="high">m0002<\/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[] = [