diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte index 2d959cfc2c..88efcc4b07 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte @@ -20,9 +20,9 @@ agenticInjectSteeringMessage } from '$lib/stores/agentic.svelte'; import { + buildSiblingInfoMap, copyToClipboard, formatMessageForClipboard, - getMessageSiblings, hasAgenticContent } from '$lib/utils'; @@ -169,6 +169,8 @@ }); }); + let siblingInfoByMessageId = $derived(buildSiblingInfoMap(allConversationMessages)); + let displayMessages = $derived.by(() => { if (!messages.length) { return []; @@ -223,18 +225,18 @@ } } - const siblingInfo = getMessageSiblings(allConversationMessages, msg.id); + const siblingInfo = siblingInfoByMessageId.get(msg.id) ?? { + message: msg, + siblingIds: [msg.id], + currentIndex: 0, + totalSiblings: 1 + }; result.push({ message: msg, toolMessages, isLastAssistantMessage: false, - siblingInfo: siblingInfo || { - message: msg, - siblingIds: [msg.id], - currentIndex: 0, - totalSiblings: 1 - } + siblingInfo }); } diff --git a/tools/ui/src/lib/utils/branching.ts b/tools/ui/src/lib/utils/branching.ts index 4e117b3c2a..c40abbdd6b 100644 --- a/tools/ui/src/lib/utils/branching.ts +++ b/tools/ui/src/lib/utils/branching.ts @@ -92,18 +92,14 @@ export function filterByLeafNodeId( * Finds the leaf node (message with no children) for a given message branch. * Traverses down the tree following the last child until reaching a leaf. * - * @param messages - All messages in the conversation + * @param nodeMap - Map of messages keyed by ID * @param messageId - Starting message ID to find leaf for * @returns The leaf node ID, or the original messageId if no children */ -export function findLeafNode(messages: readonly DatabaseMessage[], messageId: string): string { - const nodeMap = new Map(); - - // Build node map for quick lookups - for (const msg of messages) { - nodeMap.set(msg.id, msg); - } - +function findLeafNodeInMap( + nodeMap: ReadonlyMap, + messageId: string +): string { let currentNode: DatabaseMessage | undefined = nodeMap.get(messageId); while (currentNode && currentNode.children.length > 0) { // Follow the last child (most recent branch) @@ -114,6 +110,22 @@ export function findLeafNode(messages: readonly DatabaseMessage[], messageId: st return currentNode?.id ?? messageId; } +/** + * Convenience wrapper around {@link findLeafNodeInMap} for callers that only have + * a flat message array. + * + * Finds the leaf node (message with no children) for a given message branch. + * Traverses down the tree following the last child until reaching a leaf. + * + * @param messages - All messages in the conversation + * @param messageId - Starting message ID to find leaf for + * @returns The leaf node ID, or the original messageId if no children + */ +export function findLeafNode(messages: readonly DatabaseMessage[], messageId: string): string { + const nodeMap = new Map(messages.map((msg) => [msg.id, msg] as const)); + return findLeafNodeInMap(nodeMap, messageId); +} + /** * Finds all descendant messages (children, grandchildren, etc.) of a given message. * This is used for cascading deletion to remove all messages in a branch. @@ -156,21 +168,14 @@ export function findDescendantMessages( * Gets sibling information for a message, including all sibling IDs and current position. * Siblings are messages that share the same parent. * - * @param messages - All messages in the conversation + * @param nodeMap - Map of messages keyed by ID * @param messageId - The message to get sibling info for * @returns Sibling information including leaf node IDs for navigation */ export function getMessageSiblings( - messages: readonly DatabaseMessage[], + nodeMap: ReadonlyMap, messageId: string ): ChatMessageSiblingInfo | null { - const nodeMap = new Map(); - - // Build node map for quick lookups - for (const msg of messages) { - nodeMap.set(msg.id, msg); - } - const message = nodeMap.get(messageId); if (!message) { return null; @@ -203,7 +208,9 @@ export function getMessageSiblings( // Convert sibling message IDs to their corresponding leaf node IDs // This allows navigation between different conversation branches - const siblingLeafIds = siblingIds.map((siblingId: string) => findLeafNode(messages, siblingId)); + const siblingLeafIds = siblingIds.map((siblingId: string) => + findLeafNodeInMap(nodeMap, siblingId) + ); // Find current message's position among siblings const currentIndex = siblingIds.indexOf(messageId); @@ -217,85 +224,22 @@ export function getMessageSiblings( } /** - * Creates a display-ready list of messages with sibling information for UI rendering. - * This is the main function used by chat components to render conversation branches. + * Builds sibling information for every message in a conversation. + * A single node map is shared across all lookups for O(1) access. * * @param messages - All messages in the conversation - * @param leafNodeId - Current leaf node being viewed - * @returns Array of messages with sibling navigation info + * @returns Map of message ID to its sibling information */ -export function getMessageDisplayList( - messages: readonly DatabaseMessage[], - leafNodeId: string -): ChatMessageSiblingInfo[] { - // Get the current conversation path - const currentPath = filterByLeafNodeId(messages, leafNodeId, true); - const result: ChatMessageSiblingInfo[] = []; - - // Add sibling info for each message in the current path - for (const message of currentPath) { - if (message.type === 'root') { - continue; // Skip root messages in display - } - - const siblingInfo = getMessageSiblings(messages, message.id); - if (siblingInfo) { - result.push(siblingInfo); +export function buildSiblingInfoMap( + messages: readonly DatabaseMessage[] +): Map { + const nodeMap = new Map(messages.map((msg) => [msg.id, msg] as const)); + const siblingMap = new Map(); + for (const msg of messages) { + const info = getMessageSiblings(nodeMap, msg.id); + if (info) { + siblingMap.set(msg.id, info); } } - - return result; -} - -/** - * Checks if a message has multiple siblings (indicating branching at that point). - * - * @param messages - All messages in the conversation - * @param messageId - The message to check - * @returns True if the message has siblings - */ -export function hasMessageSiblings( - messages: readonly DatabaseMessage[], - messageId: string -): boolean { - const siblingInfo = getMessageSiblings(messages, messageId); - return siblingInfo ? siblingInfo.totalSiblings > 1 : false; -} - -/** - * Gets the next sibling message ID for navigation. - * - * @param messages - All messages in the conversation - * @param messageId - Current message ID - * @returns Next sibling's leaf node ID, or null if at the end - */ -export function getNextSibling( - messages: readonly DatabaseMessage[], - messageId: string -): string | null { - const siblingInfo = getMessageSiblings(messages, messageId); - if (!siblingInfo || siblingInfo.currentIndex >= siblingInfo.totalSiblings - 1) { - return null; - } - - return siblingInfo.siblingIds[siblingInfo.currentIndex + 1]; -} - -/** - * Gets the previous sibling message ID for navigation. - * - * @param messages - All messages in the conversation - * @param messageId - Current message ID - * @returns Previous sibling's leaf node ID, or null if at the beginning - */ -export function getPreviousSibling( - messages: readonly DatabaseMessage[], - messageId: string -): string | null { - const siblingInfo = getMessageSiblings(messages, messageId); - if (!siblingInfo || siblingInfo.currentIndex <= 0) { - return null; - } - - return siblingInfo.siblingIds[siblingInfo.currentIndex - 1]; + return siblingMap; } diff --git a/tools/ui/src/lib/utils/index.ts b/tools/ui/src/lib/utils/index.ts index 61b9932d3f..8474691acd 100644 --- a/tools/ui/src/lib/utils/index.ts +++ b/tools/ui/src/lib/utils/index.ts @@ -26,10 +26,7 @@ export { findLeafNode, findDescendantMessages, getMessageSiblings, - getMessageDisplayList, - hasMessageSiblings, - getNextSibling, - getPreviousSibling + buildSiblingInfoMap } from './branching'; // Code