From cdf437dc5ba44ed59c73a7c3ec691ee4b907419a Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 2 Aug 2026 07:04:52 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=AA=BA=20fix:=20Keep=20Preempt-Abandoned?= =?UTF-8?q?=20Siblings=20Nested,=20Make=20Message=20Tree=20Order-Robust=20?= =?UTF-8?q?(#14587)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ๐Ÿชบ fix: Keep Preempt-Abandoned Siblings Nested, Make Message Tree Order-Robust * ๐Ÿ”— fix: Sever Cycle Back-Edges So the Returned Message Tree Is Acyclic * ๐Ÿงช fix: Satisfy TFile in buildTree Spec fileMap Fixture * ๐ŸŒฒ fix: Assert Repaired Trees in convoStructure Specs, Uncharge Self-Parent Edges * ๐Ÿ“Œ feat: Identity-Stable Sibling Selection Across Background Tree Churn * ๐ŸŽญ test: E2E Coverage for Thread Fold and Sibling Selection Invariants * ๐Ÿ”‘ fix: Treat Newest-Sibling Re-Key as Hydration, Not a New Branch * ๐Ÿงญ fix: Rebind Sibling Selection Per Parent, Detect Appends by Membership --- .../components/Chat/Messages/MultiMessage.tsx | 111 +++++++++- .../Messages/__tests__/MultiMessage.spec.tsx | 195 ++++++++++++++++++ .../SSE/__tests__/useStepHandler.spec.ts | 57 +++++ client/src/hooks/SSE/useEventHandlers.ts | 13 +- client/src/hooks/SSE/useStepHandler.ts | 32 ++- e2e/specs/mock/thread-fold.spec.ts | 184 +++++++++++++++++ packages/data-provider/src/messages.spec.ts | 147 +++++++++++++ packages/data-provider/src/messages.ts | 67 +++++- .../src/methods/convoStructure.spec.ts | 36 +++- 9 files changed, 803 insertions(+), 39 deletions(-) create mode 100644 client/src/components/Chat/Messages/__tests__/MultiMessage.spec.tsx create mode 100644 e2e/specs/mock/thread-fold.spec.ts create mode 100644 packages/data-provider/src/messages.spec.ts diff --git a/client/src/components/Chat/Messages/MultiMessage.tsx b/client/src/components/Chat/Messages/MultiMessage.tsx index f2737d896b..dc54587aea 100644 --- a/client/src/components/Chat/Messages/MultiMessage.tsx +++ b/client/src/components/Chat/Messages/MultiMessage.tsx @@ -1,4 +1,4 @@ -import { memo, useEffect, useCallback } from 'react'; +import { memo, useEffect, useRef, useCallback } from 'react'; import { useRecoilState } from 'recoil'; import { isAssistantsEndpoint } from 'librechat-data-provider'; import type { TMessage } from 'librechat-data-provider'; @@ -9,6 +9,10 @@ import MessageParts from './MessageParts'; import Message from './Message'; import store from '~/store'; +/** First-run sentinel for `parentRef`: `messageId` itself may legitimately be + * null/undefined at the root level, so those can't mark "not yet bound". */ +const UNBOUND_PARENT: unique symbol = Symbol('multiMessageUnboundParent'); + function MultiMessage({ // messageId is used recursively here messageId, @@ -25,16 +29,107 @@ function MultiMessage({ [messagesTree?.length, setSiblingIdx], ); - useEffect(() => { - // reset siblingIdx when the tree changes, mostly when a new message is submitting. - setSiblingIdx(0); - }, [messagesTree?.length, setSiblingIdx]); + const siblingIdxRef = useRef(siblingIdx); + siblingIdxRef.current = siblingIdx; + /** Identity of this level's last committed display (`viewedId`) and its + * newest child (`newestId`), for the reconciliation below. */ + const displayedRef = useRef<{ newestId?: string; viewedId?: string }>({}); + const treeRef = useRef(null); + const parentRef = useRef(UNBOUND_PARENT); + /** + * Sibling selection is positional (reversed index), so a change to the + * children array would silently change WHAT this level displays. Reconcile + * by identity instead of blanket-resetting: + * + * - An APPENDED newest child means a submission landed here (send, + * regenerate, edit-resubmit all append) โ€” follow it, the long-standing + * behavior. An append is a newest-id change where the prior newest still + * exists; when it vanished instead, the same row was RE-KEYED (streaming + * ids hydrate to durable ids at finalize โ€” the legacy regenerate path + * mints a new UUID for `_`-suffixed ids), and following it would yank a + * user who paged away mid-stream. + * - Otherwise the change is background churn (an abandoned preempt sibling + * restored at finalize, a refetch merge dropping an optimistic row, id + * hydration) โ€” keep the message the user was viewing, recomputing its + * reversed index from its new position. Only when it no longer exists + * does the selection fall back to the newest. + * + * Keyed on tree identity via `treeRef` (streaming mints a fresh array per + * write); a plain `siblingIdx` change (the user paging the switcher) only + * records the newly viewed identity. + */ useEffect(() => { - if (messagesTree?.length && siblingIdx >= messagesTree.length) { - setSiblingIdx(0); + const length = messagesTree?.length ?? 0; + const prevTree = treeRef.current; + const treeChanged = prevTree !== messagesTree; + treeRef.current = messagesTree; + const parentChanged = parentRef.current !== messageId; + parentRef.current = messageId; + if (!messagesTree || length === 0) { + displayedRef.current = {}; + return; } - }, [siblingIdx, messagesTree?.length, setSiblingIdx]); + const newestId = messagesTree[length - 1]?.messageId; + const currentIdx = siblingIdxRef.current; + + if (parentChanged) { + /** Recursive instances are deliberately unkeyed and get REUSED across + * parents when an ancestor's branch switches: the refs still describe + * the PREVIOUS parent's children, so reconciling against them would + * wipe this parent's saved selection. Rebind to this parent's own atom + * value (clamped) instead of reconciling. */ + const boundIdx = currentIdx >= length ? 0 : currentIdx; + if (boundIdx !== currentIdx) { + setSiblingIdx(boundIdx); + } + displayedRef.current = { + newestId, + viewedId: messagesTree[length - boundIdx - 1]?.messageId, + }; + return; + } + + if (!treeChanged) { + displayedRef.current = { + newestId, + viewedId: messagesTree[length - currentIdx - 1]?.messageId, + }; + return; + } + + const previous = displayedRef.current; + /** An append means the last child is a NEW member (absent from the + * previous array) while the prior newest survived. A changed last id + * alone can also be a same-membership REORDER (sibling `createdAt` ties + * have no sort tie-breaker) or a RE-KEY (a streaming id hydrating to its + * durable id at finalize) โ€” neither is a new branch to follow. */ + const appendedNewest = + previous.newestId == null || + (newestId !== previous.newestId && + prevTree != null && + !prevTree.some((message) => message?.messageId === newestId) && + messagesTree.some((message) => message?.messageId === previous.newestId)); + let nextSiblingIdx = currentIdx; + if (appendedNewest) { + nextSiblingIdx = 0; + } else if (currentIdx > 0 && previous.viewedId != null) { + const viewedIndex = messagesTree.findIndex( + (message) => message?.messageId === previous.viewedId, + ); + nextSiblingIdx = viewedIndex >= 0 ? length - viewedIndex - 1 : 0; + } else if (currentIdx >= length) { + nextSiblingIdx = 0; + } + + if (nextSiblingIdx !== currentIdx) { + setSiblingIdx(nextSiblingIdx); + } + displayedRef.current = { + newestId, + viewedId: messagesTree[length - nextSiblingIdx - 1]?.messageId, + }; + }, [messageId, messagesTree, siblingIdx, setSiblingIdx]); if (!(messagesTree && messagesTree.length)) { return null; diff --git a/client/src/components/Chat/Messages/__tests__/MultiMessage.spec.tsx b/client/src/components/Chat/Messages/__tests__/MultiMessage.spec.tsx new file mode 100644 index 0000000000..b999d27000 --- /dev/null +++ b/client/src/components/Chat/Messages/__tests__/MultiMessage.spec.tsx @@ -0,0 +1,195 @@ +import React from 'react'; +import { RecoilRoot } from 'recoil'; +import { fireEvent, render, screen } from '@testing-library/react'; +import type { TMessage } from 'librechat-data-provider'; +import MultiMessage from '../MultiMessage'; + +type RowProps = { + message: TMessage; + siblingIdx?: number; + setSiblingIdx?: (value: number) => void; +}; + +/** Row stub exposing the sibling switcher contract (display-order index). */ +const createRowStub = () => { + const { createElement } = jest.requireActual('react'); + return ({ message, siblingIdx = 0, setSiblingIdx }: RowProps) => + createElement( + 'div', + null, + createElement('div', { 'data-testid': 'row' }, message.messageId), + createElement( + 'button', + { 'data-testid': 'prev', onClick: () => setSiblingIdx?.(siblingIdx - 1) }, + 'prev', + ), + ); +}; + +jest.mock('~/components/Messages/MessageContent', () => ({ + __esModule: true, + default: createRowStub(), +})); +jest.mock('../MessageParts', () => ({ __esModule: true, default: createRowStub() })); +jest.mock('../Message', () => ({ __esModule: true, default: createRowStub() })); + +const msg = (messageId: string): TMessage => + ({ + messageId, + parentMessageId: 'parent-1', + conversationId: 'c1', + isCreatedByUser: false, + text: messageId, + content: [{ type: 'text', text: messageId }], + children: [], + }) as unknown as TMessage; + +const tree = (ids: string[]) => ids.map((id) => msg(id)); + +const treeElement = (ids: string[]) => ( + + + +); + +const displayed = () => screen.getAllByTestId('row')[0].textContent; + +describe('MultiMessage sibling selection', () => { + it('shows the newest sibling by default and follows a newly appended one', () => { + const view = render(treeElement(['a', 'b'])); + expect(displayed()).toBe('b'); + + view.rerender(treeElement(['a', 'b', 'c'])); + expect(displayed()).toBe('c'); + }); + + /** + * Regression: background cache churn (an abandoned preempt sibling restored + * at finalize, a refetch merge) must not yank the user off the branch they + * navigated to. Only a NEW newest sibling (a submission landing at this + * level) moves the selection. + */ + it('keeps the viewed older branch when a middle sibling appears without a new newest', () => { + const view = render(treeElement(['a', 'b'])); + expect(displayed()).toBe('b'); + + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('a'); + + view.rerender(treeElement(['a', 'restored-middle', 'b'])); + expect(displayed()).toBe('a'); + }); + + it('still follows a new newest sibling from an older branch (regenerate lands)', () => { + const view = render(treeElement(['a', 'b'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('a'); + + view.rerender(treeElement(['a', 'b', 'c'])); + expect(displayed()).toBe('c'); + }); + + it('falls back to the newest when the viewed sibling disappears', () => { + const view = render(treeElement(['a', 'b', 'c'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('a'); + + view.rerender(treeElement(['b', 'c'])); + expect(displayed()).toBe('c'); + }); + + it('keeps the same-position display stable across content-only tree rebuilds', () => { + const view = render(treeElement(['a', 'b'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('a'); + + /** Streaming mints a fresh array each write with identical membership. */ + view.rerender(treeElement(['a', 'b'])); + expect(displayed()).toBe('a'); + }); + + /** + * Regression: streaming ids hydrate to durable ids at finalize (the legacy + * regenerate path mints a new UUID for `_`-suffixed preliminary ids). The + * newest child's id changing WITHOUT the prior newest surviving is a + * re-key of the same row, not an append โ€” it must not yank a user who + * paged to an older sibling mid-stream. + */ + it('keeps the viewed older branch when the newest sibling is re-keyed at finalize', () => { + const view = render(treeElement(['a', 'streaming_'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('a'); + + view.rerender(treeElement(['a', 'durable-id'])); + expect(displayed()).toBe('a'); + }); + + it('keeps the viewed branch when a re-key lands together with restored middle siblings', () => { + const view = render(treeElement(['a', 'streaming_'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('a'); + + view.rerender(treeElement(['a', 'restored-middle', 'durable-id'])); + expect(displayed()).toBe('a'); + }); + + it('still follows a genuine append after a re-key settled', () => { + const view = render(treeElement(['a', 'streaming_'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + view.rerender(treeElement(['a', 'durable-id'])); + expect(displayed()).toBe('a'); + + view.rerender(treeElement(['a', 'durable-id', 'appended'])); + expect(displayed()).toBe('appended'); + }); + + /** + * Regression: sibling `createdAt` ties have no sort tie-breaker, so a + * refetch can return the same membership in a different order. A changed + * last id without a NEW member is a reorder, not an append โ€” the viewed + * message must stay selected by identity. + */ + it('treats a same-membership reorder as churn, not an append', () => { + const view = render(treeElement(['a', 'b', 'c'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('b'); + + view.rerender(treeElement(['b', 'c', 'a'])); + expect(displayed()).toBe('b'); + }); + + /** + * Regression: the recursive instance is deliberately unkeyed and gets + * reused across parents when an ancestor's branch switches. The + * reconciliation refs then describe the previous parent's children โ€” + * reconciling against them wiped the returned-to branch's saved selection. + */ + it("preserves each parent's saved selection when the instance is reused across parents", () => { + const treeFor = (messageId: string, ids: string[]) => ( + + + + ); + + const view = render(treeFor('parent-a', ['a1', 'a2'])); + fireEvent.click(screen.getAllByTestId('prev')[0]); + expect(displayed()).toBe('a1'); + + view.rerender(treeFor('parent-b', ['b1', 'b2', 'b3'])); + expect(displayed()).toBe('b3'); + + view.rerender(treeFor('parent-a', ['a1', 'a2'])); + expect(displayed()).toBe('a1'); + }); +}); diff --git a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts index f8455d5068..170a0953ad 100644 --- a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts +++ b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts @@ -489,6 +489,63 @@ describe('useStepHandler', () => { ); }); + /** + * Regression: the preempt-fold incident. When the missing user message is + * restored while its abandoned (preempt-incomplete) sibling responses are + * already in the list, appending it at the tail orders those children + * BEFORE their parent โ€” buildTree then hoists them into phantom root + * branches and the thread folds to the latest branch until a refetch. + */ + it('inserts a missing user message before its abandoned sibling responses', () => { + const rootUser = createUserMessage({ messageId: 'user-1' }); + const assist1 = createResponseMessage({ messageId: 'assist-1', parentMessageId: 'user-1' }); + const user2 = createUserMessage({ messageId: 'user-2', parentMessageId: 'assist-1' }); + const assist2 = createResponseMessage({ messageId: 'assist-2', parentMessageId: 'user-2' }); + const abandoned1 = createResponseMessage({ + messageId: 'abandoned-1', + parentMessageId: 'user-3', + unfinished: true, + }); + const abandoned2 = createResponseMessage({ + messageId: 'abandoned-2', + parentMessageId: 'user-3', + unfinished: true, + }); + const user3 = createUserMessage({ messageId: 'user-3', parentMessageId: 'assist-2' }); + + const cachedMessages = [rootUser, assist1, user2, assist2, abandoned1, abandoned2]; + mockGetMessages.mockReturnValue(cachedMessages); + + const { result } = renderHook(() => useStepHandler(createHookParams())); + + const runStep = createRunStep({ runId: 'assist-3' }); + const submission = createSubmission({ + userMessage: user3, + messages: cachedMessages, + initialResponse: createResponseMessage({ + messageId: 'user-3_', + parentMessageId: 'user-3', + }), + }); + + act(() => { + result.current.stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, submission); + }); + + expect(mockSetMessages).toHaveBeenCalled(); + const written = mockSetMessages.mock.calls[0][0] as TMessage[]; + expect(written.map((message) => message.messageId)).toEqual([ + 'user-1', + 'assist-1', + 'user-2', + 'assist-2', + 'user-3', + 'abandoned-1', + 'abandoned-2', + 'assist-3', + ]); + }); + it('keeps the pending user message when replayed OAuth tool calls merge immediately', () => { const rootUser = createUserMessage({ messageId: 'root-user' }); const selectedResponse = createResponseMessage({ diff --git a/client/src/hooks/SSE/useEventHandlers.ts b/client/src/hooks/SSE/useEventHandlers.ts index 394c52dcd1..d8a001d7d1 100644 --- a/client/src/hooks/SSE/useEventHandlers.ts +++ b/client/src/hooks/SSE/useEventHandlers.ts @@ -450,14 +450,23 @@ export default function useEventHandlers({ (data: TSyncData, submission: EventSubmission) => { const { conversationId, thread_id, responseMessage, requestMessage } = data; const { initialResponse, messages: _messages, userMessage } = submission; - const messages = _messages.filter((msg) => msg.messageId !== userMessage.messageId); + /** Swap the optimistic user row for the server-stamped one IN PLACE. + * Filtering it out and re-appending at the tail would order any of its + * already-present children (abandoned responses from preempted + * attempts) before their parent, and the message tree hoists such rows + * into phantom root branches โ€” a folded thread. */ + const userIndex = _messages.findIndex((msg) => msg.messageId === userMessage.messageId); + const messages = + userIndex >= 0 + ? _messages.map((msg, i) => (i === userIndex ? requestMessage : msg)) + : [..._messages, requestMessage]; const nextResponseMessage = { ...initialResponse, ...responseMessage, }; - setMessages([...messages, requestMessage, nextResponseMessage]); + setMessages([...messages, nextResponseMessage]); announcePolite({ message: 'start', diff --git a/client/src/hooks/SSE/useStepHandler.ts b/client/src/hooks/SSE/useStepHandler.ts index 5b41e88373..fe3e9f0948 100644 --- a/client/src/hooks/SSE/useStepHandler.ts +++ b/client/src/hooks/SSE/useStepHandler.ts @@ -605,15 +605,28 @@ export default function useStepHandler({ return candidateMessages; } - const responseIndex = candidateMessages.findIndex( - (message) => message.messageId === responseMessageId, - ); - if (responseIndex < 0) { + /** Insert before the row's first CHILD as well as before the response + * row: abandoned responses from preempted attempts are children of + * this user message and may already sit in the list. Landing the + * parent after them orders children before their parent, which the + * message tree renders as phantom root branches (a folded thread). */ + let insertIndex = candidateMessages.length; + for (let i = 0; i < candidateMessages.length; i++) { + const message = candidateMessages[i]; + if ( + message.messageId === responseMessageId || + message.parentMessageId === userMessage.messageId + ) { + insertIndex = i; + break; + } + } + if (insertIndex >= candidateMessages.length) { return [...candidateMessages, userMessage as TMessage]; } const nextMessages = [...candidateMessages]; - nextMessages.splice(responseIndex, 0, userMessage as TMessage); + nextMessages.splice(insertIndex, 0, userMessage as TMessage); return nextMessages; }; const getResponseBaseMessages = ( @@ -790,12 +803,9 @@ export default function useStepHandler({ // Ensure userMessage is present (multi-tab: Tab 2 may not have it yet). // Regenerate reuses an existing user turn; its submission userMessage is only // a transport placeholder and must not become a new visible branch. - if ( - !submission.isRegenerate && - !updatedMessages.some((m) => m.messageId === userMessage.messageId) - ) { - updatedMessages = [...updatedMessages, userMessage as TMessage]; - } + // (`ensureUserMessagePresent` no-ops for regenerate and inserts in + // parent-before-children order otherwise.) + updatedMessages = ensureUserMessagePresent(updatedMessages, responseMessageId); setMessages([...updatedMessages, response]); } diff --git a/e2e/specs/mock/thread-fold.spec.ts b/e2e/specs/mock/thread-fold.spec.ts new file mode 100644 index 0000000000..5fc26f615f --- /dev/null +++ b/e2e/specs/mock/thread-fold.spec.ts @@ -0,0 +1,184 @@ +import { expect, test } from '@playwright/test'; +import type { Page, Response } from '@playwright/test'; +import { + MOCK_ENDPOINTS, + NEW_CHAT_PATH, + messagesView, + replyPrompt, + replyText, + selectMockEndpoint, + sendMessage, +} from './helpers'; + +/** + * Regression suite for the "folded thread" incident (PR: order-robust message + * tree + identity-stable sibling selection). The original failure: after + * preempt/interrupt churn completed a turn, the client cache held children + * ordered before their parent and the thread view collapsed to the latest + * branch (with a correct-looking sibling counter) until a reload. These tests + * pin the user-visible invariants on the real stack: every turn stays visible + * through churn, the rendered thread matches its own post-reload rendering, + * and paging to an older branch is not undone by later tree writes. + */ + +const uniqueLabel = (prefix: string) => + `${prefix}-${Date.now()}-${Math.floor(Math.random() * 1e4)}`; + +const countedPrompt = (label: string) => `E2E_COUNTED_REPLY:${label}`; +const countedReplyText = (label: string, count: number) => `E2E counted reply ${label} #${count}`; + +const messageInput = (page: Page) => page.getByRole('textbox', { name: 'Message input' }); +const messageTurns = (page: Page) => messagesView(page).locator('.message-render'); +const siblingCounter = (page: Page) => + page.getByRole('navigation', { name: 'Sibling message navigation' }).getByRole('status').first(); + +function isSteerRequest(response: Response) { + return ( + response.request().method() === 'POST' && + new URL(response.url()).pathname === '/api/agents/chat/steer' + ); +} + +async function openMockChat(page: Page) { + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); +} + +async function sendAndExpectReply(page: Page, prompt: string, reply: string) { + const response = await sendMessage(page, prompt); + expect(response.ok()).toBeTruthy(); + await expect(messagesView(page).getByText(reply)).toBeVisible({ timeout: 30000 }); +} + +async function clickSibling(page: Page, messageTextValue: string, direction: 'Previous' | 'Next') { + const render = messagesView(page) + .locator('.message-render') + .filter({ hasText: messageTextValue }) + .last(); + await render.scrollIntoViewIfNeeded(); + await render.hover(); + await render.getByRole('button', { name: `${direction} sibling message` }).click(); +} + +test.describe('thread fold regressions', () => { + test.afterEach(async ({ page }) => { + await page.evaluate(() => window.localStorage.removeItem('steerInterruptsByDefault')); + }); + + test('thread survives a mid-stream interrupt and matches its own post-reload rendering', async ({ + page, + }) => { + test.setTimeout(180000); + const label = uniqueLabel('fold-churn'); + const setupPrompt = replyPrompt(`${label}-setup`); + const setupReply = replyText(`${label}-setup`); + const interruptText = `Interrupt churn ${label}`; + + await openMockChat(page); + await sendAndExpectReply(page, setupPrompt, setupReply); + await expect(page).toHaveURL(/\/c\/[0-9a-fA-F-]{36}$/, { timeout: 15000 }); + + const run = await sendMessage(page, `E2E_SLOW_REPLY:${label}`); + expect(run.ok()).toBeTruthy(); + await expect(messagesView(page).getByText('chunk-010')).toBeVisible({ timeout: 15000 }); + + /** Queue mid-run, then escalate to an interrupt: the closest scripted + * reproduction of the incident's preempt churn (mid-stream seal, new + * generation, resume-path cache writes). */ + const input = messageInput(page); + await input.click(); + await input.fill(interruptText); + await input.press('ControlOrMeta+Enter'); + const row = page.getByTestId('queued-message-row').filter({ hasText: interruptText }); + await expect(row).toBeVisible({ timeout: 10000 }); + + const [steerResponse] = await Promise.all([ + page.waitForResponse(isSteerRequest, { timeout: 15000 }), + row.getByTestId('queued-interrupt-now').click(), + ]); + expect(steerResponse.status()).toBe(202); + + await expect( + messagesView(page).getByTestId('steer-part').filter({ hasText: interruptText }), + ).toHaveCount(1, { timeout: 90000 }); + await expect(messagesView(page).getByText(`E2E slow reply continued ${label}`)).toBeVisible({ + timeout: 30000, + }); + await expect(page.getByRole('button', { name: 'Stop generating' })).toBeHidden({ + timeout: 30000, + }); + + /** The fold's fingerprint was a live rendering that no longer matched the + * durable thread. EVERY turn must still be on screen after the churn... */ + await expect(messagesView(page).getByText(setupPrompt)).toBeVisible(); + await expect(messagesView(page).getByText(setupReply)).toBeVisible(); + await expect(messagesView(page).getByText('chunk-010')).toBeVisible(); + await expect(messageTurns(page)).toHaveCount(4); + + /** ...and reloading (the incident's only fix) must change nothing. */ + await page.reload({ timeout: 15000 }); + await expect(messagesView(page).getByText(setupPrompt)).toBeVisible({ timeout: 30000 }); + await expect(messagesView(page).getByText(setupReply)).toBeVisible(); + await expect(messagesView(page).getByText('chunk-010')).toBeVisible(); + await expect( + messagesView(page).getByTestId('steer-part').filter({ hasText: interruptText }), + ).toHaveCount(1, { timeout: 30000 }); + await expect(messageTurns(page)).toHaveCount(4); + }); + + test('older-branch selection and sibling counters survive a follow-up turn and reload', async ({ + page, + }) => { + test.setTimeout(180000); + const label = uniqueLabel('fold-branch'); + const rootPrompt = countedPrompt(label); + const firstReply = countedReplyText(label, 1); + const regeneratedReply = countedReplyText(label, 2); + const followPrompt = replyPrompt(`${label}-follow`); + const followReply = replyText(`${label}-follow`); + + await openMockChat(page); + await sendAndExpectReply(page, rootPrompt, firstReply); + await expect(page).toHaveURL(/\/c\/[0-9a-fA-F-]{36}$/, { timeout: 15000 }); + + const render = messagesView(page) + .locator('.message-render') + .filter({ hasText: firstReply }) + .last(); + await render.hover(); + await render.locator('button[title="Regenerate"]').last().click(); + await expect(messagesView(page).getByText(regeneratedReply)).toBeVisible({ timeout: 30000 }); + await expect(siblingCounter(page)).toHaveText('2 / 2'); + + /** Page to the older branch; the selection must hold, not snap back. */ + await clickSibling(page, regeneratedReply, 'Previous'); + await expect(messagesView(page).getByText(firstReply)).toBeVisible(); + await expect(messagesView(page).getByText(regeneratedReply)).toBeHidden(); + await expect(siblingCounter(page)).toHaveText('1 / 2'); + + /** A follow-up streamed from the older branch churns the tree on every + * delta and appends a deeper level โ€” none of which may move THIS level's + * selection or corrupt its counter. */ + await sendAndExpectReply(page, followPrompt, followReply); + await expect(messagesView(page).getByText(firstReply)).toBeVisible(); + await expect(messagesView(page).getByText(regeneratedReply)).toBeHidden(); + await expect(siblingCounter(page)).toHaveText('1 / 2'); + + /** Reload rebuilds selection from scratch (in-memory sibling atoms are + * gone); whichever branch the default lands on, the durable tree must be + * intact: both branches reachable through the switcher and the follow-up + * turn present on branch one. A folded tree would strand one branch. */ + await page.reload({ timeout: 15000 }); + await expect(siblingCounter(page)).toHaveText(/[12] \/ 2/, { timeout: 30000 }); + if (!(await messagesView(page).getByText(followReply).isVisible())) { + await clickSibling(page, regeneratedReply, 'Previous'); + } + await expect(messagesView(page).getByText(followReply)).toBeVisible({ timeout: 15000 }); + await expect(messagesView(page).getByText(firstReply)).toBeVisible(); + await expect(siblingCounter(page)).toHaveText('1 / 2'); + await clickSibling(page, firstReply, 'Next'); + await expect(messagesView(page).getByText(regeneratedReply)).toBeVisible(); + await expect(messagesView(page).getByText(followReply)).toBeHidden(); + await expect(siblingCounter(page)).toHaveText('2 / 2'); + }); +}); diff --git a/packages/data-provider/src/messages.spec.ts b/packages/data-provider/src/messages.spec.ts new file mode 100644 index 0000000000..1d3aed846b --- /dev/null +++ b/packages/data-provider/src/messages.spec.ts @@ -0,0 +1,147 @@ +import type { ParentMessage } from './messages'; +import type { TFile } from './types/files'; +import type { TMessage } from './types'; +import { buildTree } from './messages'; + +const msg = (messageId: string, parentMessageId: string, over: Partial = {}): TMessage => + ({ + messageId, + parentMessageId, + conversationId: 'c1', + text: '', + isCreatedByUser: false, + error: false, + ...over, + }) as TMessage; + +const asParent = (node: TMessage | undefined): ParentMessage => node as ParentMessage; + +describe('buildTree', () => { + it('returns null for null input and [] for empty input', () => { + expect(buildTree({ messages: null })).toBeNull(); + expect(buildTree({ messages: [] })).toEqual([]); + }); + + it('builds the standard tree from creation-ordered messages', () => { + const tree = buildTree({ + messages: [ + msg('u1', '00000000-0000-0000-0000-000000000000', { isCreatedByUser: true }), + msg('a1', 'u1'), + msg('u2', 'a1', { isCreatedByUser: true }), + msg('a2', 'u2'), + msg('a2b', 'u2'), + ], + }); + + expect(tree).toHaveLength(1); + const root = asParent(tree?.[0]); + expect(root.messageId).toBe('u1'); + expect(root.depth).toBe(0); + const a1 = asParent(root.children[0]); + expect(a1.messageId).toBe('a1'); + expect(a1.depth).toBe(1); + const u2 = asParent(a1.children[0]); + const siblings = u2.children.map((c) => asParent(c)); + expect(siblings.map((c) => c.messageId)).toEqual(['a2', 'a2b']); + expect(siblings.map((c) => c.siblingIndex)).toEqual([0, 1]); + expect(siblings.map((c) => c.depth)).toEqual([3, 3]); + }); + + it('nests a child that appears BEFORE its parent in the array', () => { + const tree = buildTree({ + messages: [msg('a1', 'u1'), msg('u1', '00000000-0000-0000-0000-000000000000')], + }); + + expect(tree).toHaveLength(1); + const root = asParent(tree?.[0]); + expect(root.messageId).toBe('u1'); + const a1 = asParent(root.children[0]); + expect(a1.messageId).toBe('a1'); + expect(a1.depth).toBe(1); + }); + + /** + * Regression: the preempt-fold incident. Two abandoned (preempt-incomplete) + * assistant siblings ended up ordered BEFORE their parent user message in + * the cache after a resume finalize. The single-pass builder hoisted them to + * phantom roots (while the sibling counter still read 3), folding the + * visible thread to the latest branch until a refetch restored order. + */ + it('keeps preempt-abandoned siblings nested when ordered before their parent', () => { + const NO_PARENT = '00000000-0000-0000-0000-000000000000'; + const tree = buildTree({ + messages: [ + msg('user-1', NO_PARENT, { isCreatedByUser: true }), + msg('assist-1', 'user-1'), + msg('user-2', 'assist-1', { isCreatedByUser: true }), + msg('assist-2', 'user-2', { unfinished: true }), + msg('abandoned-1', 'user-3', { unfinished: true }), + msg('abandoned-2', 'user-3', { unfinished: true }), + msg('user-3', 'assist-2', { isCreatedByUser: true }), + msg('assist-3', 'user-3'), + ], + }); + + expect(tree).toHaveLength(1); + expect(asParent(tree?.[0]).messageId).toBe('user-1'); + + const assist2 = asParent( + asParent(asParent(asParent(tree?.[0]).children[0]).children[0]).children[0], + ); + expect(assist2.messageId).toBe('assist-2'); + const user3 = asParent(assist2.children[0]); + expect(user3.messageId).toBe('user-3'); + expect(user3.depth).toBe(4); + + const siblings = user3.children.map((c) => asParent(c)); + expect(siblings.map((c) => c.messageId)).toEqual(['abandoned-1', 'abandoned-2', 'assist-3']); + expect(siblings.map((c) => c.siblingIndex)).toEqual([0, 1, 2]); + expect(siblings.map((c) => c.depth)).toEqual([5, 5, 5]); + }); + + it('surfaces a corrupt parent cycle as a root and severs the back-edge', () => { + const tree = buildTree({ + messages: [msg('u1', '00000000-0000-0000-0000-000000000000'), msg('x', 'y'), msg('y', 'x')], + }); + + const rootIds = tree?.map((node) => node.messageId); + expect(rootIds).toContain('u1'); + expect(rootIds).toContain('x'); + const cycleRoot = asParent(tree?.find((node) => node.messageId === 'x')); + const y = asParent(cycleRoot.children[0]); + expect(y.messageId).toBe('y'); + expect(y.depth).toBe(1); + /** The returned structure must be acyclic: recursive consumers + * (MultiMessage, branch utilities) walk `children` unguarded. */ + expect(y.children).toEqual([]); + }); + + it('treats a self-parented message as a root without inflating child sibling indices', () => { + const tree = buildTree({ messages: [msg('loop', 'loop'), msg('a', 'loop'), msg('b', 'loop')] }); + + expect(tree).toHaveLength(1); + const root = asParent(tree?.[0]); + expect(root.messageId).toBe('loop'); + /** The rejected self-edge must not be charged to `loop`'s child count: + * sibling indices stay within `children.length`. */ + const children = root.children.map((c) => asParent(c)); + expect(children.map((c) => c.messageId)).toEqual(['a', 'b']); + expect(children.map((c) => c.siblingIndex)).toEqual([0, 1]); + }); + + it('maps files through fileMap and skips undefined entries', () => { + const file = { file_id: 'f1', filename: 'hydrated.png' } as TFile; + const tree = buildTree({ + messages: [ + undefined, + msg('u1', '00000000-0000-0000-0000-000000000000', { + files: [{ file_id: 'f1', filename: 'stub.png' }], + }), + ], + fileMap: { f1: file }, + }); + + expect(tree).toHaveLength(1); + expect(tree?.[0].files?.[0]).toBe(file); + }); +}); diff --git a/packages/data-provider/src/messages.ts b/packages/data-provider/src/messages.ts index 2550d1ad03..d3fb4b91f6 100644 --- a/packages/data-provider/src/messages.ts +++ b/packages/data-provider/src/messages.ts @@ -2,6 +2,14 @@ import type { TFile } from './types/files'; import type { TMessage } from './types'; export type ParentMessage = TMessage & { children: TMessage[]; depth: number }; +/** + * Builds the render tree from the flat messages array. Order-robust: live + * stream/steer/preempt cache writes can momentarily place a child before its + * parent, and a single-pass link would hoist such rows into phantom root + * branches โ€” folding the visible thread to one dangling branch until a + * refetch restores creation order. Linking happens only after every message + * is indexed, so array order never changes the tree shape. + */ export function buildTree({ messages, fileMap, @@ -14,14 +22,20 @@ export function buildTree({ } const messageMap: Record = {}; - const rootMessages: TMessage[] = []; + const orderedMessages: ParentMessage[] = []; + const rootMessages: ParentMessage[] = []; const childrenCount: Record = {}; - messages.forEach((message) => { + for (const message of messages) { if (!message) { - return; + continue; } - const parentId = message.parentMessageId ?? ''; + /** A self-parented row can never link under itself (it becomes a root), + * so count it with the parentless group โ€” charging its own id would + * inflate the sibling indices of its real children past + * `children.length`. */ + const parentId = + message.parentMessageId === message.messageId ? '' : (message.parentMessageId ?? ''); childrenCount[parentId] = (childrenCount[parentId] || 0) + 1; const extendedMessage: ParentMessage = { @@ -36,15 +50,50 @@ export function buildTree({ } messageMap[message.messageId] = extendedMessage; + orderedMessages.push(extendedMessage); + } - const parentMessage = messageMap[parentId]; - if (parentMessage) { + for (const extendedMessage of orderedMessages) { + const parentMessage = messageMap[extendedMessage.parentMessageId ?? '']; + if (parentMessage && parentMessage !== extendedMessage) { parentMessage.children.push(extendedMessage); - extendedMessage.depth = parentMessage.depth + 1; } else { rootMessages.push(extendedMessage); } - }); + } - return rootMessages; + /** Depth comes from a roots-down walk (a child linked before its parent + * can't inherit depth at link time). The `visited` set doubles as the + * cycle guard: nodes on a corrupt parent cycle are unreachable from any + * root, so they resurface as roots instead of disappearing. */ + const visited = new Set(); + const assignDepths = (root: ParentMessage) => { + visited.add(root); + const stack: ParentMessage[] = [root]; + while (stack.length > 0) { + const node = stack.pop() as ParentMessage; + /** Every node has one parent, so this walk reaches each node once โ€” an + * already-visited child is a cycle back-edge. Sever it (not just skip + * it) so consumers that recurse `children` terminate. */ + if ((node.children as ParentMessage[]).some((child) => visited.has(child))) { + node.children = (node.children as ParentMessage[]).filter((child) => !visited.has(child)); + } + for (const child of node.children as ParentMessage[]) { + child.depth = node.depth + 1; + visited.add(child); + stack.push(child); + } + } + }; + for (const root of rootMessages) { + assignDepths(root); + } + for (const extendedMessage of orderedMessages) { + if (!visited.has(extendedMessage)) { + rootMessages.push(extendedMessage); + assignDepths(extendedMessage); + } + } + + return rootMessages as TMessage[]; } diff --git a/packages/data-schemas/src/methods/convoStructure.spec.ts b/packages/data-schemas/src/methods/convoStructure.spec.ts index 6ec2cf8c0e..51924b49fe 100644 --- a/packages/data-schemas/src/methods/convoStructure.spec.ts +++ b/packages/data-schemas/src/methods/convoStructure.spec.ts @@ -1,10 +1,10 @@ import mongoose from 'mongoose'; -import type { TMessage } from 'librechat-data-provider'; import { buildTree } from 'librechat-data-provider'; import { MongoMemoryServer } from 'mongodb-memory-server'; -import { createModels } from '~/models'; -import { createMessageMethods } from './message'; +import type { TMessage } from 'librechat-data-provider'; import type { IMessage } from '..'; +import { createMessageMethods } from './message'; +import { createModels } from '~/models'; jest.mock('~/config/winston', () => ({ error: jest.fn(), @@ -43,7 +43,7 @@ beforeEach(async () => { }); describe('Conversation Structure Tests', () => { - test('Conversation folding/corrupting with inconsistent timestamps', async () => { + test('Tree stays intact when inconsistent timestamps scramble retrieval order', async () => { const userId = 'testUser'; const conversationId = 'testConversation'; @@ -95,14 +95,22 @@ describe('Conversation Structure Tests', () => { // Save messages with overrideTimestamp omitted (default is false) await bulkSaveMessages(messages, true); - // Retrieve messages (this will sort by createdAt) + // Retrieve messages (this will sort by createdAt, placing message2 BEFORE its parent message1) const retrievedMessages = await getMessages({ conversationId, user: userId }); // Build tree const tree = buildTree({ messages: retrievedMessages as TMessage[] }); - // Check if the tree is incorrect (folded/corrupted) - expect(tree!.length).toBeGreaterThan(1); // Should have multiple root messages, indicating corruption + // buildTree is order-robust: a child sorted before its parent must still + // nest (the single-pass builder used to hoist it into a phantom root, + // folding the visible thread). + expect(tree!.length).toBe(1); + const root = tree![0]; + expect(root.messageId).toBe('message0'); + const message1 = root.children![0]; + expect(message1.messageId).toBe('message1'); + expect(message1.children!.map((child) => child.messageId)).toEqual(['message2', 'message3']); + expect(message1.children![0].children![0].messageId).toBe('message4'); }); test('Fix: Conversation structure maintained with more than 16 messages', async () => { @@ -139,7 +147,7 @@ describe('Conversation Structure Tests', () => { expect(currentNode.children!.length).toBe(0); // Last message should have no children }); - test('Simulate MongoDB ordering issue with more than 16 messages and close timestamps', async () => { + test('Tree stays intact under MongoDB ordering churn with close timestamps', async () => { const userId = 'testUser'; const conversationId = 'testConversation'; @@ -161,7 +169,17 @@ describe('Conversation Structure Tests', () => { await bulkSaveMessages(messages, true); const retrievedMessages = await getMessages({ conversationId, user: userId }); const tree = buildTree({ messages: retrievedMessages as TMessage[] }); - expect(tree!.length).toBeGreaterThan(1); + + // Interleaved timestamps scramble retrieval order, but the order-robust + // builder must still recover the single 20-message chain. + expect(tree!.length).toBe(1); + let currentNode = tree![0]; + for (let i = 1; i < 20; i++) { + expect(currentNode.children!.length).toBe(1); + currentNode = currentNode.children![0]; + expect(currentNode.text).toBe(`Message ${i}`); + } + expect(currentNode.children!.length).toBe(0); }); test('Fix: Preserve order with more than 16 messages by maintaining original timestamps', async () => {