Skip to content

Commit 592a5dd

Browse files
TarquinenTarquinen
committed
inject ids into all parts
Keep message references visible on every text chunk and completed tool output so multi-part messages stay fully addressable.
1 parent 425339f commit 592a5dd

File tree

4 files changed

+186
-57
lines changed

4 files changed

+186
-57
lines changed

‎lib/messages/inject/inject.ts‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,8 @@ import { compressPermission, getLastUserMessage } from "../../shared-utils"
88
import { saveSessionState } from "../../state/persistence"
99
import {
1010
appendToTextPart,
11-
appendIdToTool,
11+
appendToToolPart,
1212
createSyntheticTextPart,
13-
findLastToolPart,
1413
isIgnoredUserMessage,
1514
isProtectedUserMessage,
1615
} from "../utils"
@@ -169,7 +168,14 @@ export const injectMessageIds = (
169168
)
170169

171170
if (message.info.role === "user") {
172-
if (appendToTextPart(message, tag)) {
171+
let injected = false
172+
for (const part of message.parts) {
173+
if (part.type === "text") {
174+
injected = appendToTextPart(part, tag) || injected
175+
}
176+
}
177+
178+
if (injected) {
173179
continue
174180
}
175181

@@ -181,12 +187,16 @@ export const injectMessageIds = (
181187
continue
182188
}
183189

184-
const lastToolPart = findLastToolPart(message)
185-
if (lastToolPart && appendIdToTool(lastToolPart, tag)) {
186-
continue
190+
let injected = false
191+
for (const part of message.parts) {
192+
if (part.type === "text") {
193+
injected = appendToTextPart(part, tag) || injected
194+
} else if (part.type === "tool") {
195+
injected = appendToToolPart(part, tag) || injected
196+
}
187197
}
188198

189-
if (appendToTextPart(message, tag)) {
199+
if (injected) {
190200
continue
191201
}
192202

‎lib/messages/inject/utils.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import {
88
type MessagePriority,
99
listPriorityRefsBeforeIndex,
1010
} from "../priority"
11-
import { appendToTextPart, createSyntheticTextPart, isIgnoredUserMessage } from "../utils"
11+
import { appendToLastTextPart, createSyntheticTextPart, isIgnoredUserMessage } from "../utils"
1212
import { getLastUserMessage } from "../../shared-utils"
1313
import { getCurrentTokenUsage } from "../../strategies/utils"
1414

@@ -230,7 +230,7 @@ function injectAnchoredNudge(message: WithParts, nudgeText: string): void {
230230
return
231231
}
232232

233-
if (appendToTextPart(message, nudgeText)) {
233+
if (appendToLastTextPart(message, nudgeText)) {
234234
return
235235
}
236236

‎lib/messages/utils.ts‎

Lines changed: 22 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,24 @@ type MessagePart = WithParts["parts"][number]
7272
type ToolPart = Extract<MessagePart, { type: "tool" }>
7373
type TextPart = Extract<MessagePart, { type: "text" }>
7474

75+
export const appendToTextPart = (part: TextPart, injection: string): boolean => {
76+
if (typeof part.text !== "string") {
77+
return false
78+
}
79+
80+
const normalizedInjection = injection.replace(/^\n+/, "")
81+
if (!normalizedInjection.trim()) {
82+
return false
83+
}
84+
if (part.text.includes(normalizedInjection)) {
85+
return true
86+
}
87+
88+
const baseText = part.text.replace(/\n*$/, "")
89+
part.text = baseText.length > 0 ? `${baseText}\n\n${normalizedInjection}` : normalizedInjection
90+
return true
91+
}
92+
7593
const findLastTextPart = (message: WithParts): TextPart | null => {
7694
for (let i = message.parts.length - 1; i >= 0; i--) {
7795
const part = message.parts[i]
@@ -83,27 +101,16 @@ const findLastTextPart = (message: WithParts): TextPart | null => {
83101
return null
84102
}
85103

86-
export const appendToTextPart = (message: WithParts, injection: string): boolean => {
104+
export const appendToLastTextPart = (message: WithParts, injection: string): boolean => {
87105
const textPart = findLastTextPart(message)
88-
if (!textPart || typeof textPart.text !== "string") {
89-
return false
90-
}
91-
92-
const normalizedInjection = injection.replace(/^\n+/, "")
93-
if (!normalizedInjection.trim()) {
106+
if (!textPart) {
94107
return false
95108
}
96109

97-
const baseText = textPart.text.replace(/\n*$/, "")
98-
textPart.text =
99-
baseText.length > 0 ? `${baseText}\n\n${normalizedInjection}` : normalizedInjection
100-
return true
110+
return appendToTextPart(textPart, injection)
101111
}
102112

103-
export const appendIdToTool = (part: ToolPart, tag: string): boolean => {
104-
if (part.type !== "tool") {
105-
return false
106-
}
113+
export const appendToToolPart = (part: ToolPart, tag: string): boolean => {
107114
if (part.state?.status !== "completed" || typeof part.state.output !== "string") {
108115
return false
109116
}
@@ -115,17 +122,6 @@ export const appendIdToTool = (part: ToolPart, tag: string): boolean => {
115122
return true
116123
}
117124

118-
export const findLastToolPart = (message: WithParts): ToolPart | null => {
119-
for (let i = message.parts.length - 1; i >= 0; i--) {
120-
const part = message.parts[i]
121-
if (part.type === "tool") {
122-
return part
123-
}
124-
}
125-
126-
return null
127-
}
128-
129125
export function buildToolIdList(state: SessionState, messages: WithParts[]): string[] {
130126
const toolIds: string[] = []
131127
for (const msg of messages) {

‎tests/message-priority.test.ts‎

Lines changed: 145 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -130,10 +130,31 @@ function repeatedWord(word: string, count: number): string {
130130
return Array.from({ length: count }, () => word).join(" ")
131131
}
132132

133-
test("injectMessageIds prefers assistant tool outputs over text parts in message mode", () => {
133+
test("injectMessageIds tags every text part and tool output in message mode", () => {
134134
const sessionID = "ses_message_priority_tags"
135135
const messages: WithParts[] = [
136-
buildMessage("msg-user-1", "user", sessionID, repeatedWord("investigate", 6000), 1),
136+
{
137+
info: {
138+
id: "msg-user-1",
139+
role: "user",
140+
sessionID,
141+
agent: "assistant",
142+
model: {
143+
providerID: "anthropic",
144+
modelID: "claude-test",
145+
},
146+
time: { created: 1 },
147+
} as WithParts["info"],
148+
parts: [
149+
textPart(
150+
"msg-user-1",
151+
sessionID,
152+
"msg-user-1-part-1",
153+
repeatedWord("investigate", 6000),
154+
),
155+
textPart("msg-user-1", sessionID, "msg-user-1-part-2", "Trailing note."),
156+
],
157+
},
137158
{
138159
info: {
139160
id: "msg-assistant-1",
@@ -146,10 +167,23 @@ test("injectMessageIds prefers assistant tool outputs over text parts in message
146167
textPart(
147168
"msg-assistant-1",
148169
sessionID,
149-
"msg-assistant-1-part",
170+
"msg-assistant-1-part-1",
150171
"Short follow-up note.",
151172
),
152173
toolPart("msg-assistant-1", sessionID, "call-task-1", "task", "task output body"),
174+
textPart(
175+
"msg-assistant-1",
176+
sessionID,
177+
"msg-assistant-1-part-2",
178+
"Second text chunk.",
179+
),
180+
toolPart(
181+
"msg-assistant-1",
182+
sessionID,
183+
"call-task-2",
184+
"bash",
185+
"second tool output body",
186+
),
153187
],
154188
},
155189
]
@@ -161,31 +195,73 @@ test("injectMessageIds prefers assistant tool outputs over text parts in message
161195

162196
injectMessageIds(state, config, messages, compressionPriorities)
163197

164-
assert.equal(messages[0]?.parts.length, 1)
165-
assert.equal(messages[1]?.parts.length, 2)
166-
167-
const userText = messages[0]?.parts[0]
168-
const assistantText = messages[1]?.parts[0]
169-
const assistantTool = messages[1]?.parts[1]
170-
171-
assert.equal(userText?.type, "text")
172-
assert.equal(assistantText?.type, "text")
173-
assert.equal(assistantTool?.type, "tool")
198+
assert.equal(messages[0]?.parts.length, 2)
199+
assert.equal(messages[1]?.parts.length, 4)
200+
201+
const userTextOne = messages[0]?.parts[0]
202+
const userTextTwo = messages[0]?.parts[1]
203+
const assistantTextOne = messages[1]?.parts[0]
204+
const assistantToolOne = messages[1]?.parts[1]
205+
const assistantTextTwo = messages[1]?.parts[2]
206+
const assistantToolTwo = messages[1]?.parts[3]
207+
208+
assert.equal(userTextOne?.type, "text")
209+
assert.equal(userTextTwo?.type, "text")
210+
assert.equal(assistantTextOne?.type, "text")
211+
assert.equal(assistantToolOne?.type, "tool")
212+
assert.equal(assistantTextTwo?.type, "text")
213+
assert.equal(assistantToolTwo?.type, "tool")
174214
assert.match(
175-
(userText as any).text,
215+
(userTextOne as any).text,
176216
/\n\n<dcp-message-id priority="high">m0001<\/dcp-message-id>/,
177217
)
178-
assert.equal((assistantText as any).text, "Short follow-up note.")
179218
assert.match(
180-
(assistantTool as any).state.output,
219+
(userTextTwo as any).text,
220+
/\n\n<dcp-message-id priority="high">m0001<\/dcp-message-id>/,
221+
)
222+
assert.match(
223+
(assistantTextOne as any).text,
224+
/\n\n<dcp-message-id priority="low">m0002<\/dcp-message-id>/,
225+
)
226+
assert.match(
227+
(assistantToolOne as any).state.output,
228+
/<dcp-message-id priority="low">m0002<\/dcp-message-id>/,
229+
)
230+
assert.match(
231+
(assistantTextTwo as any).text,
232+
/\n\n<dcp-message-id priority="low">m0002<\/dcp-message-id>/,
233+
)
234+
assert.match(
235+
(assistantToolTwo as any).state.output,
181236
/<dcp-message-id priority="low">m0002<\/dcp-message-id>/,
182237
)
183238
})
184239

185-
test("injectMessageIds marks protected user messages as BLOCKED without priority in message mode", () => {
240+
test("injectMessageIds marks every protected user text part as BLOCKED in message mode", () => {
186241
const sessionID = "ses_message_blocked_user_tags"
187242
const messages: WithParts[] = [
188-
buildMessage("msg-user-1", "user", sessionID, repeatedWord("investigate", 6000), 1),
243+
{
244+
info: {
245+
id: "msg-user-1",
246+
role: "user",
247+
sessionID,
248+
agent: "assistant",
249+
model: {
250+
providerID: "anthropic",
251+
modelID: "claude-test",
252+
},
253+
time: { created: 1 },
254+
} as WithParts["info"],
255+
parts: [
256+
textPart(
257+
"msg-user-1",
258+
sessionID,
259+
"msg-user-1-part-1",
260+
repeatedWord("investigate", 6000),
261+
),
262+
textPart("msg-user-1", sessionID, "msg-user-1-part-2", "Trailing note."),
263+
],
264+
},
189265
buildMessage("msg-assistant-1", "assistant", sessionID, "Short follow-up note.", 2),
190266
]
191267
const state = createSessionState()
@@ -197,19 +273,66 @@ test("injectMessageIds marks protected user messages as BLOCKED without priority
197273

198274
injectMessageIds(state, config, messages, compressionPriorities)
199275

200-
const userText = messages[0]?.parts[0]
276+
const userTextOne = messages[0]?.parts[0]
277+
const userTextTwo = messages[0]?.parts[1]
201278
const assistantText = messages[1]?.parts[0]
202279

203-
assert.equal(userText?.type, "text")
280+
assert.equal(userTextOne?.type, "text")
281+
assert.equal(userTextTwo?.type, "text")
204282
assert.equal(assistantText?.type, "text")
205-
assert.match((userText as any).text, /\n\n<dcp-message-id>BLOCKED<\/dcp-message-id>/)
206-
assert.doesNotMatch((userText as any).text, /priority=/)
283+
assert.match((userTextOne as any).text, /\n\n<dcp-message-id>BLOCKED<\/dcp-message-id>/)
284+
assert.match((userTextTwo as any).text, /\n\n<dcp-message-id>BLOCKED<\/dcp-message-id>/)
285+
assert.doesNotMatch((userTextOne as any).text, /priority=/)
286+
assert.doesNotMatch((userTextTwo as any).text, /priority=/)
207287
assert.match(
208288
(assistantText as any).text,
209289
/\n\n<dcp-message-id priority="low">m0002<\/dcp-message-id>/,
210290
)
211291
})
212292

293+
test("injectMessageIds tags every text part and tool output in range mode", () => {
294+
const sessionID = "ses_range_message_id_tags"
295+
const messages: WithParts[] = [
296+
buildMessage("msg-user-1", "user", sessionID, repeatedWord("investigate", 6000), 1),
297+
{
298+
info: {
299+
id: "msg-assistant-1",
300+
role: "assistant",
301+
sessionID,
302+
agent: "assistant",
303+
time: { created: 2 },
304+
} as WithParts["info"],
305+
parts: [
306+
textPart("msg-assistant-1", sessionID, "msg-assistant-1-part-1", "First chunk."),
307+
toolPart("msg-assistant-1", sessionID, "call-task-range-1", "task", "first output"),
308+
textPart("msg-assistant-1", sessionID, "msg-assistant-1-part-2", "Second chunk."),
309+
toolPart(
310+
"msg-assistant-1",
311+
sessionID,
312+
"call-task-range-2",
313+
"bash",
314+
"second output",
315+
),
316+
],
317+
},
318+
]
319+
const state = createSessionState()
320+
const config = buildConfig("range")
321+
322+
assignMessageRefs(state, messages)
323+
injectMessageIds(state, config, messages)
324+
325+
const assistantTextOne = messages[1]?.parts[0]
326+
const assistantToolOne = messages[1]?.parts[1]
327+
const assistantTextTwo = messages[1]?.parts[2]
328+
const assistantToolTwo = messages[1]?.parts[3]
329+
330+
assert.match((assistantTextOne as any).text, /\n\n<dcp-message-id>m0002<\/dcp-message-id>/)
331+
assert.match((assistantToolOne as any).state.output, /<dcp-message-id>m0002<\/dcp-message-id>/)
332+
assert.match((assistantTextTwo as any).text, /\n\n<dcp-message-id>m0002<\/dcp-message-id>/)
333+
assert.match((assistantToolTwo as any).state.output, /<dcp-message-id>m0002<\/dcp-message-id>/)
334+
})
335+
213336
test("message-mode nudges append to existing text parts and list only earlier visible high-priority message IDs", () => {
214337
const sessionID = "ses_message_priority_nudges"
215338
const messages: WithParts[] = [

0 commit comments

Comments
 (0)