Skip to content

Commit 5767be1

Browse files
TarquinenTarquinen
committed
fix: tighten assistant injection guards
1 parent 4f35959 commit 5767be1

4 files changed

Lines changed: 191 additions & 14 deletions

File tree

‎lib/messages/inject/inject.ts‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
appendToTextPart,
1111
appendToToolPart,
1212
createSyntheticTextPart,
13+
hasContent,
1314
isIgnoredUserMessage,
1415
isProtectedUserMessage,
1516
} from "../utils"
@@ -186,15 +187,7 @@ export const injectMessageIds = (
186187
continue
187188
}
188189

189-
const hasContent = message.parts.some(
190-
(p) =>
191-
(p.type === "text" && typeof p.text === "string" && p.text.trim().length > 0) ||
192-
(p.type === "tool" &&
193-
p.state?.status === "completed" &&
194-
typeof p.state.output === "string"),
195-
)
196-
197-
if (!hasContent) {
190+
if (!hasContent(message)) {
198191
continue
199192
}
200193

‎lib/messages/inject/utils.ts‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,13 @@ import {
88
type MessagePriority,
99
listPriorityRefsBeforeIndex,
1010
} from "../priority"
11-
import { appendToLastTextPart, createSyntheticTextPart, isIgnoredUserMessage } from "../utils"
11+
import {
12+
appendToTextPart,
13+
appendToLastTextPart,
14+
createSyntheticTextPart,
15+
hasContent,
16+
isIgnoredUserMessage,
17+
} from "../utils"
1218
import { getLastUserMessage } from "../../shared-utils"
1319
import { getCurrentTokenUsage } from "../../strategies/utils"
1420
import { getActiveSummaryTokenUsage } from "../../state/utils"
@@ -236,11 +242,11 @@ function injectAnchoredNudge(message: WithParts, nudgeText: string): void {
236242
return
237243
}
238244

239-
if (appendToLastTextPart(message, nudgeText)) {
240-
return
241-
}
242-
243245
if (message.info.role === "user") {
246+
if (appendToLastTextPart(message, nudgeText)) {
247+
return
248+
}
249+
244250
message.parts.push(createSyntheticTextPart(message, nudgeText))
245251
return
246252
}
@@ -249,6 +255,18 @@ function injectAnchoredNudge(message: WithParts, nudgeText: string): void {
249255
return
250256
}
251257

258+
for (const part of message.parts) {
259+
if (part.type === "text") {
260+
if (appendToTextPart(part, nudgeText)) {
261+
return
262+
}
263+
}
264+
}
265+
266+
if (!hasContent(message)) {
267+
return
268+
}
269+
252270
const syntheticPart = createSyntheticTextPart(message, nudgeText)
253271
const firstToolIndex = message.parts.findIndex((p) => p.type === "tool")
254272
if (firstToolIndex === -1) {

‎lib/messages/utils.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,18 @@ export const appendToLastTextPart = (message: WithParts, injection: string): boo
109109
return appendToTextPart(textPart, injection)
110110
}
111111

112+
export const hasContent = (message: WithParts): boolean => {
113+
return message.parts.some(
114+
(part) =>
115+
(part.type === "text" &&
116+
typeof part.text === "string" &&
117+
part.text.trim().length > 0) ||
118+
(part.type === "tool" &&
119+
part.state?.status === "completed" &&
120+
typeof part.state.output === "string"),
121+
)
122+
}
123+
112124
export const appendToToolPart = (part: ToolPart, tag: string): boolean => {
113125
if (part.state?.status !== "completed" || typeof part.state.output !== "string") {
114126
return false

‎tests/message-priority.test.ts‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -511,6 +511,160 @@ test("range-mode nudges append to existing text parts before tool outputs", () =
511511
assert.equal((toolOutput as any).state.output, "task output body")
512512
})
513513

514+
test("range-mode nudges inject only once for assistant messages with multiple text parts", () => {
515+
const sessionID = "ses_range_nudge_multi_text"
516+
const messages: WithParts[] = [
517+
buildMessage("msg-user-1", "user", sessionID, "Hello", 1),
518+
{
519+
info: {
520+
id: "msg-assistant-1",
521+
role: "assistant",
522+
sessionID,
523+
agent: "assistant",
524+
time: { created: 2 },
525+
} as WithParts["info"],
526+
parts: [
527+
textPart("msg-assistant-1", sessionID, "assistant-text-1", "First chunk."),
528+
textPart("msg-assistant-1", sessionID, "assistant-text-2", "Second chunk."),
529+
],
530+
},
531+
]
532+
const state = createSessionState()
533+
const config = buildConfig("range")
534+
535+
assignMessageRefs(state, messages)
536+
state.nudges.contextLimitAnchors.add("msg-assistant-1")
537+
538+
applyAnchoredNudges(state, config, messages, {
539+
system: "",
540+
compressRange: "",
541+
compressMessage: "",
542+
contextLimitNudge: "<dcp-system-reminder>Base context nudge</dcp-system-reminder>",
543+
turnNudge: "<dcp-system-reminder>Base turn nudge</dcp-system-reminder>",
544+
iterationNudge: "<dcp-system-reminder>Base iteration nudge</dcp-system-reminder>",
545+
})
546+
547+
assert.match((messages[1]?.parts[0] as any).text, /Base context nudge/)
548+
assert.doesNotMatch((messages[1]?.parts[1] as any).text, /Base context nudge/)
549+
})
550+
551+
test("range-mode nudges skip empty assistant messages to avoid prefill (issue #463)", () => {
552+
const sessionID = "ses_range_nudge_empty_assistant"
553+
const messages: WithParts[] = [
554+
buildMessage("msg-user-1", "user", sessionID, "Hello", 1),
555+
{
556+
info: {
557+
id: "msg-assistant-empty",
558+
role: "assistant",
559+
sessionID,
560+
agent: "assistant",
561+
time: { created: 2 },
562+
} as WithParts["info"],
563+
parts: [],
564+
},
565+
]
566+
const state = createSessionState()
567+
const config = buildConfig("range")
568+
569+
assignMessageRefs(state, messages)
570+
state.nudges.contextLimitAnchors.add("msg-assistant-empty")
571+
572+
applyAnchoredNudges(state, config, messages, {
573+
system: "",
574+
compressRange: "",
575+
compressMessage: "",
576+
contextLimitNudge: "<dcp-system-reminder>Base context nudge</dcp-system-reminder>",
577+
turnNudge: "<dcp-system-reminder>Base turn nudge</dcp-system-reminder>",
578+
iterationNudge: "<dcp-system-reminder>Base iteration nudge</dcp-system-reminder>",
579+
})
580+
581+
assert.equal(messages[1]?.parts.length, 0)
582+
})
583+
584+
test("range-mode nudges skip assistant with only pending tool parts (issue #463)", () => {
585+
const sessionID = "ses_range_nudge_pending_tool"
586+
const messages: WithParts[] = [
587+
buildMessage("msg-user-1", "user", sessionID, "Hello", 1),
588+
{
589+
info: {
590+
id: "msg-assistant-pending",
591+
role: "assistant",
592+
sessionID,
593+
agent: "assistant",
594+
time: { created: 2 },
595+
} as WithParts["info"],
596+
parts: [
597+
{
598+
id: "pending-tool-part",
599+
messageID: "msg-assistant-pending",
600+
sessionID,
601+
type: "tool" as const,
602+
tool: "bash",
603+
callID: "call-pending-1",
604+
state: {
605+
status: "pending" as const,
606+
input: { command: "ls" },
607+
},
608+
} as any,
609+
],
610+
},
611+
]
612+
const state = createSessionState()
613+
const config = buildConfig("range")
614+
615+
assignMessageRefs(state, messages)
616+
state.nudges.contextLimitAnchors.add("msg-assistant-pending")
617+
618+
applyAnchoredNudges(state, config, messages, {
619+
system: "",
620+
compressRange: "",
621+
compressMessage: "",
622+
contextLimitNudge: "<dcp-system-reminder>Base context nudge</dcp-system-reminder>",
623+
turnNudge: "<dcp-system-reminder>Base turn nudge</dcp-system-reminder>",
624+
iterationNudge: "<dcp-system-reminder>Base iteration nudge</dcp-system-reminder>",
625+
})
626+
627+
assert.equal(messages[1]?.parts.length, 1)
628+
assert.equal(messages[1]?.parts[0]?.type, "tool")
629+
})
630+
631+
test("range-mode nudges append to an assistant empty text part (issue #463)", () => {
632+
const sessionID = "ses_range_nudge_empty_text"
633+
const messages: WithParts[] = [
634+
buildMessage("msg-user-1", "user", sessionID, "Hello", 1),
635+
{
636+
info: {
637+
id: "msg-assistant-empty-text",
638+
role: "assistant",
639+
sessionID,
640+
agent: "assistant",
641+
time: { created: 2 },
642+
} as WithParts["info"],
643+
parts: [textPart("msg-assistant-empty-text", sessionID, "empty-text-part", "")],
644+
},
645+
]
646+
const state = createSessionState()
647+
const config = buildConfig("range")
648+
649+
assignMessageRefs(state, messages)
650+
state.nudges.contextLimitAnchors.add("msg-assistant-empty-text")
651+
652+
applyAnchoredNudges(state, config, messages, {
653+
system: "",
654+
compressRange: "",
655+
compressMessage: "",
656+
contextLimitNudge: "<dcp-system-reminder>Base context nudge</dcp-system-reminder>",
657+
turnNudge: "<dcp-system-reminder>Base turn nudge</dcp-system-reminder>",
658+
iterationNudge: "<dcp-system-reminder>Base iteration nudge</dcp-system-reminder>",
659+
})
660+
661+
assert.equal(messages[1]?.parts.length, 1)
662+
assert.match(
663+
(messages[1]?.parts[0] as any).text,
664+
/<dcp-system-reminder>Base context nudge[\s\S]*Compressed block context:/,
665+
)
666+
})
667+
514668
test("message-mode rendered compressed summaries mark block IDs as BLOCKED", () => {
515669
const sessionID = "ses_message_blocked_blocks"
516670
const messages: WithParts[] = [

0 commit comments

Comments
 (0)