append injections to text parts

Author: Tarquinen

lib/messages/inject/inject.ts

@@ -7,6 +7,7 @@ import type { CompressionPriorityMap } from "../priority"
 import { compressPermission, getLastUserMessage } from "../../shared-utils"
 import { saveSessionState } from "../../state/persistence"
 import {
+    appendToTextPart,
     appendIdToTool,
     createSyntheticTextPart,
     findLastToolPart,
@@ -162,6 +163,10 @@ export const injectMessageIds = (
                 : undefined
         const tag = formatMessageIdTag(messageRef, priority ? { priority } : undefined)
 
+        if (appendToTextPart(message, tag)) {
+            continue
+        }
+
         if (message.info.role === "user") {
             message.parts.push(createSyntheticTextPart(message, tag))
             continue

lib/messages/inject/utils.ts

@@ -8,7 +8,7 @@ import {
     type MessagePriority,
     listPriorityRefsBeforeIndex,
 } from "../priority"
-import { createSyntheticTextPart, isIgnoredUserMessage } from "../utils"
+import { appendToTextPart, createSyntheticTextPart, isIgnoredUserMessage } from "../utils"
 import { getLastUserMessage } from "../../shared-utils"
 import { getCurrentTokenUsage } from "../../strategies/utils"
 
@@ -192,20 +192,20 @@ export function buildCompressedBlockGuidance(state: SessionState): string {
     ].join("\n")
 }
 
-function appendGuidanceToDcpTag(hintText: string, guidance: string): string {
+function appendGuidanceToDcpTag(nudgeText: string, guidance: string): string {
     if (!guidance.trim()) {
-        return hintText
+        return nudgeText
     }
 
     const closeTag = "</dcp-system-reminder>"
-    const closeTagIndex = hintText.lastIndexOf(closeTag)
+    const closeTagIndex = nudgeText.lastIndexOf(closeTag)
 
     if (closeTagIndex === -1) {
-        return hintText
+        return nudgeText
     }
 
-    const beforeClose = hintText.slice(0, closeTagIndex).trimEnd()
-    const afterClose = hintText.slice(closeTagIndex)
+    const beforeClose = nudgeText.slice(0, closeTagIndex).trimEnd()
+    const afterClose = nudgeText.slice(closeTagIndex)
     return `${beforeClose}\n\n${guidance}\n${afterClose}`
 }
 
@@ -225,21 +225,25 @@ function buildMessagePriorityGuidance(
     return renderMessagePriorityGuidance(priorityLabel, refs)
 }
 
-function injectAnchoredNudge(message: WithParts, hintText: string): void {
-    if (!hintText.trim()) {
+function injectAnchoredNudge(message: WithParts, nudgeText: string): void {
+    if (!nudgeText.trim()) {
+        return
+    }
+
+    if (appendToTextPart(message, nudgeText)) {
         return
     }
 
     if (message.info.role === "user") {
-        message.parts.push(createSyntheticTextPart(message, hintText))
+        message.parts.push(createSyntheticTextPart(message, nudgeText))
         return
     }
 
     if (message.info.role !== "assistant") {
         return
     }
 
-    const syntheticPart = createSyntheticTextPart(message, hintText)
+    const syntheticPart = createSyntheticTextPart(message, nudgeText)
     const firstToolIndex = message.parts.findIndex((p) => p.type === "tool")
     if (firstToolIndex === -1) {
         message.parts.push(syntheticPart)
@@ -291,23 +295,23 @@ function collectTurnNudgeAnchors(
 function applyRangeModeAnchoredNudge(
     anchorMessageIds: Set<string>,
     messages: WithParts[],
-    basePrompt: string,
+    baseNudgeText: string,
     compressedBlockGuidance: string,
 ): void {
-    const hintText = appendGuidanceToDcpTag(basePrompt, compressedBlockGuidance)
-    if (!hintText.trim()) {
+    const nudgeText = appendGuidanceToDcpTag(baseNudgeText, compressedBlockGuidance)
+    if (!nudgeText.trim()) {
         return
     }
 
     for (const { message } of collectAnchoredMessages(anchorMessageIds, messages)) {
-        injectAnchoredNudge(message, hintText)
+        injectAnchoredNudge(message, nudgeText)
     }
 }
 
 function applyMessageModeAnchoredNudge(
     anchorMessageIds: Set<string>,
     messages: WithParts[],
-    basePrompt: string,
+    baseNudgeText: string,
     compressionPriorities?: CompressionPriorityMap,
 ): void {
     for (const { message, index } of collectAnchoredMessages(anchorMessageIds, messages)) {
@@ -317,8 +321,8 @@ function applyMessageModeAnchoredNudge(
             index,
             MESSAGE_MODE_NUDGE_PRIORITY,
         )
-        const hintText = appendGuidanceToDcpTag(basePrompt, priorityGuidance)
-        injectAnchoredNudge(message, hintText)
+        const nudgeText = appendGuidanceToDcpTag(baseNudgeText, priorityGuidance)
+        injectAnchoredNudge(message, nudgeText)
     }
 }
 

lib/messages/utils.ts

@@ -67,6 +67,35 @@ export const createSyntheticTextPart = (
 
 type MessagePart = WithParts["parts"][number]
 type ToolPart = Extract<MessagePart, { type: "tool" }>
+type TextPart = Extract<MessagePart, { type: "text" }>
+
+const findLastTextPart = (message: WithParts): TextPart | null => {
+    for (let i = message.parts.length - 1; i >= 0; i--) {
+        const part = message.parts[i]
+        if (part.type === "text") {
+            return part
+        }
+    }
+
+    return null
+}
+
+export const appendToTextPart = (message: WithParts, injection: string): boolean => {
+    const textPart = findLastTextPart(message)
+    if (!textPart || typeof textPart.text !== "string") {
+        return false
+    }
+
+    const normalizedInjection = injection.replace(/^\n+/, "")
+    if (!normalizedInjection.trim()) {
+        return false
+    }
+
+    const baseText = textPart.text.replace(/\n*$/, "")
+    textPart.text =
+        baseText.length > 0 ? `${baseText}\n\n${normalizedInjection}` : normalizedInjection
+    return true
+}
 
 export const appendIdToTool = (part: ToolPart, tag: string): boolean => {
     if (part.type !== "tool") {

tests/message-priority.test.ts

@@ -1,15 +1,15 @@
 import assert from "node:assert/strict"
 import test from "node:test"
+import type { PluginConfig } from "../lib/config"
+import { createTextCompleteHandler } from "../lib/hooks"
 import { assignMessageRefs } from "../lib/message-ids"
-import { buildPriorityMap } from "../lib/messages/priority"
 import { injectMessageIds } from "../lib/messages/inject/inject"
 import { applyAnchoredNudges } from "../lib/messages/inject/utils"
+import { buildPriorityMap } from "../lib/messages/priority"
 import { stripHallucinationsFromString } from "../lib/messages/utils"
 import { createSessionState, type WithParts } from "../lib/state"
-import type { PluginConfig } from "../lib/config"
-import { createTextCompleteHandler } from "../lib/hooks"
 
-function buildConfig(): PluginConfig {
+function buildConfig(mode: "message" | "range" = "message"): PluginConfig {
     return {
         enabled: true,
         debug: false,
@@ -33,7 +33,7 @@ function buildConfig(): PluginConfig {
         },
         protectedFilePatterns: [],
         compress: {
-            mode: "message",
+            mode,
             permission: "allow",
             showCompression: false,
             maxContextLimit: 150000,
@@ -68,6 +68,28 @@ function textPart(messageID: string, sessionID: string, id: string, text: string
     }
 }
 
+function toolPart(
+    messageID: string,
+    sessionID: string,
+    callID: string,
+    toolName: string,
+    output: string,
+) {
+    return {
+        id: `${callID}-part`,
+        messageID,
+        sessionID,
+        type: "tool" as const,
+        tool: toolName,
+        callID,
+        state: {
+            status: "completed" as const,
+            input: { description: "demo" },
+            output,
+        },
+    }
+}
+
 function buildMessage(
     id: string,
     role: "user" | "assistant",
@@ -106,11 +128,28 @@ function repeatedWord(word: string, count: number): string {
     return Array.from({ length: count }, () => word).join(" ")
 }
 
-test("injectMessageIds adds priority attributes in message mode", () => {
+test("injectMessageIds appends priority tags to existing text parts in message mode", () => {
     const sessionID = "ses_message_priority_tags"
     const messages: WithParts[] = [
         buildMessage("msg-user-1", "user", sessionID, repeatedWord("investigate", 6000), 1),
-        buildMessage("msg-assistant-1", "assistant", sessionID, "Short follow-up note.", 2),
+        {
+            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",
+                    "Short follow-up note.",
+                ),
+                toolPart("msg-assistant-1", sessionID, "call-task-1", "task", "task output body"),
+            ],
+        },
     ]
     const state = createSessionState()
     const config = buildConfig()
@@ -120,19 +159,28 @@ test("injectMessageIds adds priority attributes in message mode", () => {
 
     injectMessageIds(state, config, messages, compressionPriorities)
 
-    const userTag = messages[0]?.parts[messages[0].parts.length - 1]
-    const assistantTag = messages[1]?.parts[messages[1].parts.length - 1]
+    assert.equal(messages[0]?.parts.length, 1)
+    assert.equal(messages[1]?.parts.length, 2)
 
-    assert.equal(userTag?.type, "text")
-    assert.equal(assistantTag?.type, "text")
-    assert.match((userTag as any).text, /<dcp-message-id priority="high">m0001<\/dcp-message-id>/)
+    const userText = messages[0]?.parts[0]
+    const assistantText = messages[1]?.parts[0]
+    const assistantTool = messages[1]?.parts[1]
+
+    assert.equal(userText?.type, "text")
+    assert.equal(assistantText?.type, "text")
+    assert.equal(assistantTool?.type, "tool")
     assert.match(
-        (assistantTag as any).text,
-        /<dcp-message-id priority="low">m0002<\/dcp-message-id>/,
+        (userText as any).text,
+        /\n\n<dcp-message-id priority="high">m0001<\/dcp-message-id>/,
     )
+    assert.match(
+        (assistantText as any).text,
+        /\n\n<dcp-message-id priority="low">m0002<\/dcp-message-id>/,
+    )
+    assert.equal((assistantTool as any).state.output, "task output body")
 })
 
-test("message-mode nudges list only earlier visible high-priority message IDs", () => {
+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[] = [
         buildMessage("msg-user-1", "user", sessionID, repeatedWord("alpha", 6000), 1),
@@ -168,15 +216,64 @@ test("message-mode nudges list only earlier visible high-priority message IDs",
         compressionPriorities,
     )
 
-    const injectedNudge = messages[2]?.parts[messages[2].parts.length - 1]
+    assert.equal(messages[2]?.parts.length, 1)
+
+    const injectedNudge = messages[2]?.parts[0]
     assert.equal(injectedNudge?.type, "text")
+    assert.match((injectedNudge as any).text, /\n\n<dcp-system-reminder>Base context nudge/)
     assert.match((injectedNudge as any).text, /Message priority context:/)
     assert.match((injectedNudge as any).text, /High-priority message IDs before this point: m0001/)
     assert.doesNotMatch((injectedNudge as any).text, /m0002/)
     assert.doesNotMatch((injectedNudge as any).text, /m0003/)
     assert.doesNotMatch((injectedNudge as any).text, /m0004/)
 })
 
+test("range-mode nudges append to existing text parts before tool outputs", () => {
+    const sessionID = "ses_range_nudge_injection"
+    const messages: WithParts[] = [
+        buildMessage("msg-user-1", "user", sessionID, repeatedWord("alpha", 6000), 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", "Working summary."),
+                toolPart("msg-assistant-1", sessionID, "call-task-2", "task", "task output body"),
+            ],
+        },
+    ]
+    const state = createSessionState()
+    const config = buildConfig("range")
+
+    assignMessageRefs(state, messages)
+    state.prune.messages.activeBlockIds.add(7)
+    state.nudges.contextLimitAnchors.add("msg-assistant-1")
+
+    applyAnchoredNudges(state, config, messages, {
+        system: "",
+        compressRange: "",
+        compressMessage: "",
+        contextLimitNudge: "<dcp-system-reminder>Base context nudge</dcp-system-reminder>",
+        turnNudge: "<dcp-system-reminder>Base turn nudge</dcp-system-reminder>",
+        iterationNudge: "<dcp-system-reminder>Base iteration nudge</dcp-system-reminder>",
+    })
+
+    assert.equal(messages[1]?.parts.length, 2)
+
+    const injectedNudge = messages[1]?.parts[0]
+    const toolOutput = messages[1]?.parts[1]
+    assert.equal(injectedNudge?.type, "text")
+    assert.equal(toolOutput?.type, "tool")
+    assert.match((injectedNudge as any).text, /\n\n<dcp-system-reminder>Base context nudge/)
+    assert.match((injectedNudge as any).text, /Compressed block context:/)
+    assert.match((injectedNudge as any).text, /Active compressed blocks in this session: 1 \(b7\)/)
+    assert.equal((toolOutput as any).state.output, "task output body")
+})
+
 test("hallucination stripping removes exact metadata tags and preserves lookalikes", async () => {
     const text =
         'alpha<dcp-message-id priority="high">m0007</dcp-message-id>' +