Skip to content

Commit 441947e

Browse files
committed
Drop tool results orphaned by trimMessagesToFitTokenLimit
The removal run in trimMessagesToFitTokenLimit stops as soon as the token budget is met, which can land between an assistant tool-call and its role:'tool' result. The surviving tool message then reaches the provider without its call and the step fails with 'tool_call_id does not exist' — observed on the find-files request path via getMessagesSubset. After the removal loop, drop results whose call this trim removed; results whose call never existed in the input history pass through unchanged so pre-existing orphans are not silently rewritten, and providerExecuted calls are excluded to mirror the pairing semantics of dropUnansweredToolCalls.
1 parent 1310581 commit 441947e

2 files changed

Lines changed: 207 additions & 0 deletions

File tree

packages/agent-runtime/src/util/__tests__/messages.test.ts

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -511,6 +511,175 @@ describe('trimMessagesToFitTokenLimit', () => {
511511
expect(replacementMessages.length).toBeGreaterThan(0)
512512
})
513513
})
514+
515+
describe('orphaned tool results at the removal boundary', () => {
516+
// Regression: the removal run stops as soon as the token budget is met,
517+
// which can land exactly between an assistant tool-call and its result.
518+
// The kept tool message then reaches the provider without its call and
519+
// is rejected with "tool_call_id does not exist", failing the whole step.
520+
521+
const toolCallPart = (toolCallId: string, filler: string) => ({
522+
type: 'tool-call' as const,
523+
toolCallId,
524+
toolName: 'write_file' as const,
525+
input: { content: filler },
526+
})
527+
528+
const toolResultMessage = (
529+
toolCallId: string,
530+
filler: string,
531+
): Message => ({
532+
role: 'tool',
533+
toolName: 'write_file',
534+
toolCallId,
535+
content: jsonToolResult(filler),
536+
})
537+
538+
/** Every tool result in the output must have its call in the output. */
539+
const expectNoOrphanedToolResults = (result: Message[]) => {
540+
const keptToolCallIds = new Set<string>()
541+
for (const message of result) {
542+
if (message.role !== 'assistant' || !Array.isArray(message.content)) {
543+
continue
544+
}
545+
for (const part of message.content) {
546+
if (part.type === 'tool-call') {
547+
keptToolCallIds.add(part.toolCallId)
548+
}
549+
}
550+
}
551+
const orphaned = result.filter(
552+
(message) =>
553+
message.role === 'tool' && !keptToolCallIds.has(message.toolCallId),
554+
)
555+
expect(orphaned).toEqual([])
556+
}
557+
558+
it('drops a tool result whose call was removed at the boundary', () => {
559+
const messages: Message[] = [
560+
userMessage('please write the file'),
561+
assistantMessage({
562+
content: [toolCallPart('c1', 'x'.repeat(4000))],
563+
}),
564+
toolResultMessage('c1', 'ok'),
565+
assistantMessage('done'),
566+
]
567+
568+
const result = trimMessagesToFitTokenLimit({
569+
messages,
570+
systemTokens: 0,
571+
maxTotalTokens: 600,
572+
logger,
573+
})
574+
575+
expectNoOrphanedToolResults(result)
576+
// The final 'done' assistant message must survive the trim.
577+
expect(
578+
result.some(
579+
(message) =>
580+
message.role === 'assistant' &&
581+
message.content.some(
582+
(part) => part.type === 'text' && part.text === 'done',
583+
),
584+
),
585+
).toBe(true)
586+
})
587+
588+
it('drops tool results orphaned after a kept keepDuringTruncation message', () => {
589+
// The removal run is not a pure prefix when keepDuringTruncation
590+
// messages sit in the middle: the boundary orphan can appear anywhere,
591+
// not just leading the kept run.
592+
const messages: Message[] = [
593+
userMessage('please write the file'),
594+
assistantMessage({
595+
content: [toolCallPart('c1', 'x'.repeat(4000))],
596+
}),
597+
userMessage({ content: 'steer', keepDuringTruncation: true }),
598+
toolResultMessage('c1', 'ok'),
599+
assistantMessage('done'),
600+
]
601+
602+
const result = trimMessagesToFitTokenLimit({
603+
messages,
604+
systemTokens: 0,
605+
maxTotalTokens: 600,
606+
logger,
607+
})
608+
609+
expectNoOrphanedToolResults(result)
610+
})
611+
612+
it('keeps tool results whose call survives the trim', () => {
613+
const messages: Message[] = [
614+
userMessage('please write the file'),
615+
assistantMessage({
616+
content: [toolCallPart('c1', 'x'.repeat(4000))],
617+
}),
618+
toolResultMessage('c1', 'ok'),
619+
assistantMessage('done'),
620+
]
621+
622+
// Generous budget: nothing gets removed, pairing stays intact.
623+
const result = trimMessagesToFitTokenLimit({
624+
messages,
625+
systemTokens: 0,
626+
maxTotalTokens: 60_000,
627+
logger,
628+
})
629+
630+
expect(result).toEqual(messages)
631+
})
632+
633+
it('leaves pre-existing orphans in an already-malformed history untouched', () => {
634+
// The tool result has no call anywhere in the input: trimming did not
635+
// create this orphan, so this fix leaves it alone instead of silently
636+
// rewriting histories it did not break. The trim is active here (the
637+
// large user message is removed) — only calls removed BY the trim
638+
// cause their results to be dropped.
639+
const messages: Message[] = [
640+
userMessage('x'.repeat(4000)),
641+
assistantMessage('done'),
642+
toolResultMessage('ghost', 'ok'),
643+
]
644+
645+
const result = trimMessagesToFitTokenLimit({
646+
messages,
647+
systemTokens: 0,
648+
maxTotalTokens: 600,
649+
logger,
650+
})
651+
652+
const ghost = result.find(
653+
(message) => message.role === 'tool' && message.toolCallId === 'ghost',
654+
)
655+
expect(ghost).toBeDefined()
656+
})
657+
658+
it('keeps the invariant across a sweep of budgets', () => {
659+
const messages: Message[] = [
660+
userMessage('please write the file'),
661+
assistantMessage({
662+
content: [toolCallPart('c1', 'x'.repeat(3000))],
663+
}),
664+
toolResultMessage('c1', 'ok'),
665+
assistantMessage({
666+
content: [toolCallPart('c2', 'y'.repeat(1500))],
667+
}),
668+
toolResultMessage('c2', 'ok'),
669+
assistantMessage('done'),
670+
]
671+
672+
for (const maxTotalTokens of [200, 400, 800, 1600, 3200, 6400]) {
673+
const result = trimMessagesToFitTokenLimit({
674+
messages,
675+
systemTokens: 0,
676+
maxTotalTokens,
677+
logger,
678+
})
679+
expectNoOrphanedToolResults(result)
680+
}
681+
})
682+
})
514683
})
515684

516685
describe('getPreviouslyReadFiles', () => {

packages/agent-runtime/src/util/messages.ts

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,44 @@ export function trimMessagesToFitTokenLimit(params: {
229229
}
230230
}
231231

232+
// The removal run stops as soon as the token budget is met, which can land
233+
// between an assistant tool-call and its result: the surviving tool message
234+
// then reaches the provider without its call and is rejected with
235+
// "tool_call_id does not exist", failing the whole step. Drop results whose
236+
// call this trim removed — but only those, so orphans that were already in
237+
// the input history pass through unchanged.
238+
const removedCallIds = new Set<string>()
239+
const survivingCallIds = new Set<string>()
240+
const collectCallIds = (message: Message, into: Set<string>) => {
241+
if (message.role !== 'assistant' || !Array.isArray(message.content)) {
242+
return
243+
}
244+
for (const part of message.content) {
245+
if (part.type === 'tool-call' && part.providerExecuted !== true) {
246+
into.add(part.toolCallId)
247+
}
248+
}
249+
}
250+
for (const message of messages) {
251+
collectCallIds(message, removedCallIds)
252+
}
253+
if (removedCallIds.size > 0) {
254+
for (const message of filteredMessages) {
255+
if (message === placeholder) continue
256+
collectCallIds(message, survivingCallIds)
257+
}
258+
for (let i = filteredMessages.length - 1; i >= 0; i--) {
259+
const message = filteredMessages[i]
260+
if (message === placeholder || message.role !== 'tool') continue
261+
if (
262+
removedCallIds.has(message.toolCallId) &&
263+
!survivingCallIds.has(message.toolCallId)
264+
) {
265+
filteredMessages.splice(i, 1)
266+
}
267+
}
268+
}
269+
232270
return filteredMessages.map((m) =>
233271
m === placeholder ? replacementMessage : m,
234272
)

0 commit comments

Comments
 (0)