mirror of
https://github.com/lobehub/lobe-chat.git
synced 2026-06-13 19:20:04 +00:00
🐛 fix(chat): preserve message order after tool results (#15657)
This commit is contained in:
@@ -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;
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user