diff --git a/packages/conversation-flow/src/__tests__/parse.test.ts b/packages/conversation-flow/src/__tests__/parse.test.ts index c340e2fa4b..11101d0806 100644 --- a/packages/conversation-flow/src/__tests__/parse.test.ts +++ b/packages/conversation-flow/src/__tests__/parse.test.ts @@ -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) { @@ -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', () => { diff --git a/packages/conversation-flow/src/transformation/FlatListBuilder.ts b/packages/conversation-flow/src/transformation/FlatListBuilder.ts index 92b47594cf..c90ab835c3 100644 --- a/packages/conversation-flow/src/transformation/FlatListBuilder.ts +++ b/packages/conversation-flow/src/transformation/FlatListBuilder.ts @@ -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, + 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, + 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, + ): { 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): 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 */ diff --git a/src/store/chat/slices/message/selectors/displayMessage.test.ts b/src/store/chat/slices/message/selectors/displayMessage.test.ts index 784338b52e..10d165a886 100644 --- a/src/store/chat/slices/message/selectors/displayMessage.test.ts +++ b/src/store/chat/slices/message/selectors/displayMessage.test.ts @@ -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=, 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 = { + 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', () => { diff --git a/src/store/chat/slices/message/selectors/displayMessage.ts b/src/store/chat/slices/message/selectors/displayMessage.ts index ce79b6fd44..eb7b0c19f6 100644 --- a/src/store/chat/slices/message/selectors/displayMessage.ts +++ b/src/store/chat/slices/message/selectors/displayMessage.ts @@ -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; };