🐛 fix(chat): preserve message order after tool results (#15657)

This commit is contained in:
Neko
2026-06-11 16:18:18 +08:00
committed by GitHub
parent 14501ea69a
commit a60d11df48
4 changed files with 343 additions and 65 deletions
@@ -1,6 +1,7 @@
import { describe, expect, it } from 'vitest';
import { parse } from '../parse';
import type { Message } from '../types/shared';
import { inputs, outputs } from './fixtures';
function serializeParseResult(result: ReturnType<typeof parse>) {
@@ -161,6 +162,194 @@ describe('parse', () => {
]);
expect((currentGroup as any).children[0].tools[0].result_msg_id).toBe('tool-current-1');
});
it('should keep sibling assistant continuations before later user turns under another tool result', () => {
const time = (seconds: number) =>
new Date(`2026-01-01T00:00:${String(seconds).padStart(2, '0')}.000Z`).getTime();
const messages: Message[] = [
{ content: 'root', createdAt: time(0), id: 'u0', role: 'user', updatedAt: time(0) },
{
content: 'assistant with two tools',
createdAt: time(1),
id: 'a0',
parentId: 'u0',
role: 'assistant',
tools: [
{
apiName: 'update',
arguments: '{}',
id: 'tc-later',
identifier: 'internal',
result_msg_id: 'tool-later',
type: 'default',
},
{
apiName: 'read',
arguments: '{}',
id: 'tc-first',
identifier: 'internal',
result_msg_id: 'tool-first',
type: 'default',
},
],
updatedAt: time(1),
},
{
content: 'todo updated',
createdAt: time(2),
id: 'tool-later',
parentId: 'a0',
role: 'tool',
tool_call_id: 'tc-later',
updatedAt: time(2),
},
{
content: 'context result',
createdAt: time(3),
id: 'tool-first',
parentId: 'a0',
role: 'tool',
tool_call_id: 'tc-first',
updatedAt: time(3),
},
{
content: 'summary before question',
createdAt: time(4),
id: 'summary',
parentId: 'tool-first',
role: 'assistant',
updatedAt: time(4),
},
{
content: 'Earlier assistant continuation',
createdAt: time(5),
id: 'first-question',
parentId: 'tool-first',
role: 'assistant',
updatedAt: time(5),
},
{
content: 'later answer',
createdAt: time(6),
id: 'later-answer',
parentId: 'tool-later',
role: 'assistant',
updatedAt: time(6),
},
{
content: 'Later user follow-up',
createdAt: time(7),
id: 'status-user',
parentId: 'tool-later',
role: 'user',
updatedAt: time(7),
},
{
content: '...',
createdAt: time(8),
id: 'status-assistant',
parentId: 'status-user',
role: 'assistant',
updatedAt: time(8),
},
];
const result = parse(messages);
const ids = result.flatList.map((message) => message.id);
// ROOT CAUSE:
//
// If one assistantGroup owns multiple tool results, and one tool result has
// multiple assistant children while another tool result later receives user
// turns, FlatListBuilder currently walks the other tool result first. That
// renders later user turns before the earlier assistant continuation.
//
// We fixed this by preserving the chronological continuation order across
// sibling tool-result children instead of deferring the second assistant.
expect(ids.indexOf('first-question')).toBeLessThan(ids.indexOf('status-user'));
});
it('should interleave continuations from sibling tool results by child creation time', () => {
const time = (seconds: number) =>
new Date(`2026-01-01T00:01:${String(seconds).padStart(2, '0')}.000Z`).getTime();
const messages: Message[] = [
{ content: 'root', createdAt: time(0), id: 'u0', role: 'user', updatedAt: time(0) },
{
content: 'assistant with sibling tools',
createdAt: time(1),
id: 'a0',
parentId: 'u0',
role: 'assistant',
tools: [
{
apiName: 'first',
arguments: '{}',
id: 'tc-a',
identifier: 'internal',
result_msg_id: 'tool-a',
type: 'default',
},
{
apiName: 'second',
arguments: '{}',
id: 'tc-b',
identifier: 'internal',
result_msg_id: 'tool-b',
type: 'default',
},
],
updatedAt: time(1),
},
{
content: 'tool a result',
createdAt: time(2),
id: 'tool-a',
parentId: 'a0',
role: 'tool',
tool_call_id: 'tc-a',
updatedAt: time(2),
},
{
content: 'tool b result',
createdAt: time(3),
id: 'tool-b',
parentId: 'a0',
role: 'tool',
tool_call_id: 'tc-b',
updatedAt: time(3),
},
{
content: 'first user continuation',
createdAt: time(4),
id: 'tool-a-first',
parentId: 'tool-a',
role: 'user',
updatedAt: time(4),
},
{
content: 'middle user continuation',
createdAt: time(5),
id: 'tool-b-middle',
parentId: 'tool-b',
role: 'user',
updatedAt: time(5),
},
{
content: 'last user continuation',
createdAt: time(6),
id: 'tool-a-last',
parentId: 'tool-a',
role: 'user',
updatedAt: time(6),
},
];
const result = parse(messages);
const ids = result.flatList.map((message) => message.id);
expect(ids.indexOf('tool-a-first')).toBeLessThan(ids.indexOf('tool-b-middle'));
expect(ids.indexOf('tool-b-middle')).toBeLessThan(ids.indexOf('tool-a-last'));
});
});
describe('Branching', () => {
@@ -262,29 +262,13 @@ export class FlatListBuilder {
processedIds.add(completion.id);
}
// Continue after the assistant chain
// Priority 1: If last assistant has non-tool children, continue from it
// Priority 2: Otherwise continue from tools (for cases where user replies to tool)
const lastAssistant = assistantChain.at(-1);
const toolIds = new Set(allToolMessages.map((t) => t.id));
const lastAssistantNonToolChildren = lastAssistant
? this.childrenMap.get(lastAssistant.id)?.filter((childId) => !toolIds.has(childId))
: undefined;
if (
lastAssistantNonToolChildren &&
lastAssistantNonToolChildren.length > 0 &&
lastAssistant
) {
// Follow-up messages exist after the last assistant (not tools)
this.buildFlatListRecursive(lastAssistant.id, flatList, processedIds, allMessages);
} else {
// No non-tool children of last assistant, check tools for children
for (const toolMsg of allToolMessages) {
this.buildFlatListRecursive(toolMsg.id, flatList, processedIds, allMessages);
}
}
this.continueAfterAssistantGroup(
assistantChain,
allToolMessages,
flatList,
processedIds,
allMessages,
);
continue;
}
@@ -409,25 +393,13 @@ export class FlatListBuilder {
assistantChain.forEach((m) => processedIds.add(m.id));
allToolMessages.forEach((m) => processedIds.add(m.id));
// Continue after the assistant chain
const lastAssistant = assistantChain.at(-1);
const toolIds = new Set(allToolMessages.map((t) => t.id));
const lastAssistantNonToolChildren = lastAssistant
? this.childrenMap.get(lastAssistant.id)?.filter((childId) => !toolIds.has(childId))
: undefined;
if (
lastAssistantNonToolChildren &&
lastAssistantNonToolChildren.length > 0 &&
lastAssistant
) {
this.buildFlatListRecursive(lastAssistant.id, flatList, processedIds, allMessages);
} else {
for (const toolMsg of allToolMessages) {
this.buildFlatListRecursive(toolMsg.id, flatList, processedIds, allMessages);
}
}
this.continueAfterAssistantGroup(
assistantChain,
allToolMessages,
flatList,
processedIds,
allMessages,
);
} else {
// Regular assistant message (not assistantGroup) - add branch info
const activeBranchWithBranches = this.createMessageWithBranches(
@@ -528,27 +500,96 @@ export class FlatListBuilder {
assistantChain.forEach((m) => processedIds.add(m.id));
allToolMessages.forEach((m) => processedIds.add(m.id));
// Continue after the assistant chain
// Priority 1: If last assistant has non-tool children, continue from it
// Priority 2: Otherwise continue from tools (for cases where user replies to tool)
this.continueAfterAssistantGroup(
assistantChain,
allToolMessages,
flatList,
processedIds,
allMessages,
);
}
private continueAfterAssistantGroup(
assistantChain: Message[],
allToolMessages: Message[],
flatList: Message[],
processedIds: Set<string>,
allMessages: Message[],
): void {
const lastAssistant = assistantChain.at(-1);
const toolIds = new Set(allToolMessages.map((t) => t.id));
const parentIds = [
...(lastAssistant ? [lastAssistant.id] : []),
...allToolMessages.map((toolMessage) => toolMessage.id),
];
const lastAssistantNonToolChildren = lastAssistant
? this.childrenMap.get(lastAssistant.id)?.filter((childId) => !toolIds.has(childId))
: undefined;
while (true) {
const nextContinuation = this.findNextUnprocessedChild(parentIds, processedIds);
if (!nextContinuation) return;
if (lastAssistantNonToolChildren && lastAssistantNonToolChildren.length > 0 && lastAssistant) {
// Follow-up messages exist after the last assistant (not tools)
this.buildFlatListRecursive(lastAssistant.id, flatList, processedIds, allMessages);
} else {
// No non-tool children of last assistant, check tools for children
for (const toolMsg of allToolMessages) {
this.buildFlatListRecursive(toolMsg.id, flatList, processedIds, allMessages);
if (this.shouldDrainParentContinuations(nextContinuation.parentId, processedIds)) {
this.buildFlatListRecursive(
nextContinuation.parentId,
flatList,
processedIds,
allMessages,
);
continue;
}
this.buildFlatListRecursiveForChild(
nextContinuation.parentId,
nextContinuation.child.id,
flatList,
processedIds,
allMessages,
);
}
}
private buildFlatListRecursiveForChild(
parentId: string,
childId: string,
flatList: Message[],
processedIds: Set<string>,
allMessages: Message[],
): void {
const childIds = this.childrenMap.get(parentId) ?? [];
this.childrenMap.set(parentId, [childId]);
try {
this.buildFlatListRecursive(parentId, flatList, processedIds, allMessages);
} finally {
this.childrenMap.set(parentId, childIds);
}
}
private findNextUnprocessedChild(
parentIds: string[],
processedIds: Set<string>,
): { child: Message; parentId: string } | undefined {
return parentIds
.flatMap((parentId) =>
(this.childrenMap.get(parentId) ?? [])
.map((childId) => this.messageMap.get(childId))
.filter((child): child is Message => !!child && !processedIds.has(child.id))
.map((child) => ({ child, parentId })),
)
.sort((a, b) => a.child.createdAt - b.child.createdAt)[0];
}
private shouldDrainParentContinuations(parentId: string, processedIds: Set<string>): boolean {
const parentMessage = this.messageMap.get(parentId);
const children = (this.childrenMap.get(parentId) ?? []).filter(
(childId) => !processedIds.has(childId),
);
if (!parentMessage || children.length <= 1) return false;
if (this.isAgentCouncilMode(parentMessage)) return true;
const taskChildren = children.filter((childId) => this.messageMap.get(childId)?.role === 'task');
return taskChildren.length > 1;
}
/**
* Check if message has compare mode in metadata
*/
@@ -794,7 +794,7 @@ describe('displayMessageSelectors', () => {
expect(result).toBeUndefined();
});
it('should return last child with tools result_msg_id', () => {
it('should return last child id even when the child has tools', () => {
const messageWithChildrenAndTools = {
id: 'msg-1',
role: 'assistantGroup',
@@ -829,7 +829,62 @@ describe('displayMessageSelectors', () => {
};
const result = displayMessageSelectors.findLastMessageId('msg-1')(state as ChatStore);
expect(result).toBe('tool-result-id');
expect(result).toBe('child-2');
});
it('should not use an internal tool result as the parent for a new user message', () => {
// ROOT CAUSE:
//
// When the visible last message is an assistantGroup and its last block has
// tools, findLastMessageId returns the tool result id. sendMessage then
// stores the new user message with parentId=<tool-result-id>, which makes
// normal user turns siblings under a tool message instead of continuing
// the conversational chain.
//
// Before patch:
// assistantGroup -> last block -> internal tool result
// new user message parentId = tool result id
//
// Expected behavior:
// new user message parentId = the visible conversational block id
const assistantGroup = {
id: 'assistant-group-1',
role: 'assistantGroup',
content: '',
children: [
{
id: 'assistant-block-1',
content: 'I will run an internal tool.',
},
{
id: 'assistant-block-2',
content: 'Internal tool completed.',
tools: [
{
id: 'internal-tool-call',
identifier: 'test-runner',
apiName: 'internal_update',
arguments: '{}',
type: 'default',
result_msg_id: 'internal-tool-result-message',
},
],
},
],
} as unknown as UIChatMessage;
const state: Partial<ChatStore> = {
activeAgentId: 'test-id',
messagesMap: {
[messageMapKey({ agentId: 'test-id' })]: [assistantGroup],
},
};
const result = displayMessageSelectors.findLastMessageId('assistant-group-1')(
state as ChatStore,
);
expect(result).toBe('assistant-block-2');
});
it('should return lastMessageId for compressedGroup instead of group id', () => {
@@ -235,13 +235,6 @@ const getGroupLatestMessageWithoutTools = (id: string) => (s: ChatStoreState) =>
const findLastBlockId = (block: AssistantContentBlock | undefined): string | undefined => {
if (!block) return undefined;
// Check tools for result message ID
if (block.tools && block.tools.length > 0) {
const lastTool = block.tools.at(-1);
return lastTool?.result_msg_id;
}
// Return block ID
return block.id;
};