From 9618be6eb3824ea4269ad29ddf1e2f17a7ddbb4e Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 14 Jun 2026 09:38:06 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=8C=BF=20fix:=20Preserve=20Viewed=20Branc?= =?UTF-8?q?h=20on=20Sibling-Tree=20Churn=20(#13732)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🌿 fix: Preserve Viewed Branch on Sibling-Tree Churn Regenerating a message could snap the view to an unrelated newest branch. MultiMessage reset siblingIdx to 0 (newest) on any messagesTree.length change, but getRegenerateSubmissionMessages slices the flat message array during a regenerate — the streaming handlers render a tree missing unrelated sibling branches, then finalHandler restores the full set. That 2→1→2 child-count swing snapped unrelated forks to their newest sibling, so regenerating the latest response on an older branch jumped to a previously regenerated branch. Replace the indiscriminate reset with per-fork branch memory: a 'seen' set distinguishes a genuinely new sibling (submission/regeneration/edit here — focus it) from one transiently dropped and restored (preserve the user's branch). Decision extracted as the pure, unit-tested resolveSiblingSelection. - client/src/utils/messages.ts: resolveSiblingSelection + tests - MultiMessage: seen/selectedId refs, structural id-signature effect - e2e: regenerate-latest-on-older-branch keeps the viewed branch (fails on the old reset, passes now) * 🧪 test: Long-Thread Branch Preservation E2E Add the user-reported scenario: in a multi-turn thread, regenerate an earlier response (forking a root branch), switch back to the original, then regenerate a later response on it — the original branch must stay intact. Uses labeled prompts so each turn's unique reply is a reliable settle signal. Verified it fails on the original MultiMessage and passes with the fix. * 🎨 style: Fix import order in MultiMessage (react before recoil) * 🌿 fix: Keep Unrelated Branches in Regenerate Optimistic Render Regenerating a message used a flat `messages.slice(0, targetIndex)` for the optimistic render, which also drops unrelated sibling branches that merely sit later in the flat array. Mid-regenerate the thread briefly collapsed to a short branch (visible flash) and the scroll jumped to the shrunken content and didn't recover — the same flat-array root cause as the branch-reset bug. Remove only the regenerated response and its descendants, keeping unrelated branches. The thread (and scroll) stay put through the regenerate. This array is render-only — the server regenerates from parentMessageId and createPayload doesn't include it — so summing by subtree never affects the request. Verified via a small-viewport scroll trace: old collapses 903->295px / 8->2 renders mid-stream; fixed stays 903px / 8 renders, scroll held at bottom. Unit test covers the keep-unrelated-branches behavior (fails on the old slice). * 🌿 fix: Let an Explicit Branch Selection Survive Streaming ID Churn resolveSiblingSelection focused any unseen sibling id before checking the committed selection. When an in-flight response's id is replaced mid-stream (placeholder → server/run id, e.g. useStepHandler re-keys to runId) after the user switched to a different sibling, that swap looked like a brand-new sibling and stole focus back to the streaming branch. Reorder: the committed selection wins while still present; only focus a fresh sibling when the selection is gone (regenerated away, or its own placeholder id was just replaced — that's how a regen/edit still takes focus, since the slice removes the old response). Added unit tests for both churn directions. * 🌿 fix: Only Focus a New Sibling When the Fork Actually Grew The previous churn fix (selection-wins-first) was too aggressive: a genuinely new sibling ADDED while the prior selection is still present — e.g. a follow-up re-parented as a sibling after a generation-start failure — was no longer focused, so its reply never rendered (broke message-tree generation-start recovery e2e). Gate new-sibling focus on actual growth: resolveSiblingSelection now takes prevCount and only focuses a never-seen id when ids.length > prevCount. A same-count placeholder→server id swap (churn) or a restored already-seen sibling is not growth, so the committed selection still wins there. Covers follow-up/new-branch focus, churn steal-prevention, and self-churn follow. message-tree + chat e2e: 17 passed (incl. the recovered generation-start test). * 🌿 refactor: Drop MultiMessage Branch-Memory in Favor of the Slice Fix The regenerate-slice fix (keep unrelated branches in the optimistic render) is the true root cause: with no spurious tree collapse, the original setSiblingIdx(0)-on-length-change never misfires, so the branch-reset is fixed without per-fork memory. The earlier MultiMessage rewrite (seen/selectedId/ prevCount + resolveSiblingSelection) was a symptom patch added before the root cause was found, and its per-instance memory generated two edge-case findings (placeholder→server id churn; divergence from external siblingIdx writes like resume restore). Revert MultiMessage to the simple upstream version and remove resolveSiblingSelection (+ its tests). The slice fix + the existing branch e2e (chat.spec: switch-back, regenerate-latest, long-thread) cover the behavior; all 17 chat + message-tree branch specs pass with this version. * 🌿 fix: Focus the Regenerated Response When Its Fork Count Is Unchanged When a parent already has multiple sibling responses and the user switches to a non-latest one and regenerates it, the optimistic slice drops the target but keeps the other siblings, so the child count is unchanged. MultiMessage only resets the (reversed) sibling index on a length change, so the stale index kept pointing at the kept sibling and the regenerating response stayed hidden until the server restored the dropped sibling at finalize (count bump → reset). Explicitly focus the newest sibling (reversed index 0 = the appended response) of the regenerated fork in createdHandler. Position-based, fires only on the regenerate action, so it doesn't reintroduce the placeholder→server id churn or external-write fragility that a per-render selection memory had. E2E: new during-stream test (slow+counted reply marker) asserting the regenerating response is visible before finalize; negatively verified (fails without the focus call, passes with it). * 🌿 fix: Eliminate Pre-Created Flash by Focusing at the Optimistic Render The createdHandler focus removed the until-finalize bug, but a brief flash remained between clicking regenerate and the `created` event: useChatFunctions renders the optimistic placeholder first, and that render has the same unchanged-count problem, so the kept sibling showed until createdHandler fired. Extract the focus into a shared useFocusRegeneratedResponse hook and apply it at the optimistic render too (useChatFunctions) and on `created` (useEventHandlers). The placeholder is now focused from the first frame. E2E: gated pre-created test — holds the SSE stream GET (the chat POST returns a stream id; the stream is a separate GET) so `created` cannot arrive, leaving only the optimistic render, then asserts the kept sibling is already gone. This isolates the optimistic focus (createdHandler cannot mask it); negatively verified (fails without the optimistic focus call). * 🧪 test: Extend Store Mock for the Regenerate Focus Hook useChatFunctions.regenerate.spec.tsx mocks ~/store and recoil partially; the new useFocusRegeneratedResponse calls store.messagesSiblingIdxFamily via a recoil `set`, neither of which the mock provided (TypeError on regenerate). Add messagesSiblingIdxFamily to the store mock and `set` to the useRecoilCallback mock. Test-only; production code unchanged. --- .../useChatFunctions.regenerate.spec.tsx | 2 + .../Chat/__tests__/useChatFunctions.spec.ts | 36 ++++++ client/src/hooks/Chat/useChatFunctions.ts | 33 ++++- .../hooks/Chat/useFocusRegeneratedResponse.ts | 25 ++++ client/src/hooks/SSE/useEventHandlers.ts | 5 + e2e/setup/fake-model.js | 16 +++ e2e/specs/mock/chat.spec.ts | 120 ++++++++++++++++++ e2e/specs/mock/message-tree.spec.ts | 101 +++++++++++++++ 8 files changed, 333 insertions(+), 5 deletions(-) create mode 100644 client/src/hooks/Chat/useFocusRegeneratedResponse.ts diff --git a/client/src/hooks/Chat/__tests__/useChatFunctions.regenerate.spec.tsx b/client/src/hooks/Chat/__tests__/useChatFunctions.regenerate.spec.tsx index 9e7cb76b85..ab5112f0e6 100644 --- a/client/src/hooks/Chat/__tests__/useChatFunctions.regenerate.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useChatFunctions.regenerate.spec.tsx @@ -31,6 +31,7 @@ jest.mock('recoil', () => ({ snapshot: { getLoadable: () => ({ state: 'hasValue', contents: [] }), }, + set: jest.fn(), reset: jest.fn(), }), })); @@ -48,6 +49,7 @@ jest.mock('~/store', () => ({ isSubmittingFamily: () => 'isSubmitting', showStopButtonByIndex: () => 'showStopButton', pendingManualSkillsByConvoId: () => 'pendingManualSkills', + messagesSiblingIdxFamily: () => 'messagesSiblingIdx', }, useGetEphemeralAgent: () => mockGetEphemeralAgent, })); diff --git a/client/src/hooks/Chat/__tests__/useChatFunctions.spec.ts b/client/src/hooks/Chat/__tests__/useChatFunctions.spec.ts index ae80a48dcc..0835cb43f4 100644 --- a/client/src/hooks/Chat/__tests__/useChatFunctions.spec.ts +++ b/client/src/hooks/Chat/__tests__/useChatFunctions.spec.ts @@ -90,4 +90,40 @@ describe('regenerate response targeting', () => { }).map((message) => message.messageId), ).toEqual(['user-1']); }); + + it('keeps unrelated sibling branches when regenerating (no flat-array drop)', () => { + // user-1 has two responses: the original chain (assistant-1 -> ... ) and a + // regenerated sibling (assistant-1b) that sits LATER in the flat array. + const messages = [ + userMessage('user-1'), + assistantMessage('assistant-1', 'user-1'), + userMessage('user-2', 'assistant-1'), + assistantMessage('assistant-2', 'user-2'), + assistantMessage('assistant-1b', 'user-1'), + ]; + + // Regenerating the latest response on the original branch must drop only + // that response, NOT the unrelated assistant-1b branch. + expect( + getRegenerateSubmissionMessages({ + messages, + targetResponseMessage: messages[3], + initialResponseId: 'assistant-2_', + }) + .map((message) => message.messageId) + .sort(), + ).toEqual(['assistant-1', 'assistant-1b', 'user-1', 'user-2']); + + // Regenerating an earlier response drops its subtree (user-2, assistant-2) + // but still keeps the unrelated assistant-1b branch. + expect( + getRegenerateSubmissionMessages({ + messages, + targetResponseMessage: messages[1], + initialResponseId: 'assistant-1_', + }) + .map((message) => message.messageId) + .sort(), + ).toEqual(['assistant-1b', 'user-1']); + }); }); diff --git a/client/src/hooks/Chat/useChatFunctions.ts b/client/src/hooks/Chat/useChatFunctions.ts index 95df4f933b..e7fda4758f 100644 --- a/client/src/hooks/Chat/useChatFunctions.ts +++ b/client/src/hooks/Chat/useChatFunctions.ts @@ -32,6 +32,7 @@ import { createDualMessageContent, getRouteChatProjectId, } from '~/utils'; +import useFocusRegeneratedResponse from '~/hooks/Chat/useFocusRegeneratedResponse'; import useSetFilesToDelete from '~/hooks/Files/useSetFilesToDelete'; import useGetSender from '~/hooks/Conversations/useGetSender'; import store, { useGetEphemeralAgent } from '~/store'; @@ -149,12 +150,32 @@ export function getRegenerateSubmissionMessages({ initialResponseId?: string | null; }): TMessage[] { if (targetResponseMessage?.messageId) { - const targetIndex = messages.findIndex( - (message) => message.messageId === targetResponseMessage.messageId, - ); - if (targetIndex >= 0) { - return messages.slice(0, targetIndex); + /** + * Remove the response being regenerated and its descendants only — NOT a + * flat `slice(0, targetIndex)`, which also drops unrelated sibling branches + * that merely sit later in the array. That collapse made the optimistic + * render briefly lose other branches mid-regenerate (visible flash, and the + * scroll jumping to the shrunken content). Keeping them holds the thread — + * and scroll — steady. This array is render-only; the server regenerates + * from `parentMessageId`, so removing by subtree never affects the payload. + */ + const removed = new Set([targetResponseMessage.messageId]); + let grew = true; + while (grew) { + grew = false; + for (const message of messages) { + const parentMessageId = message.parentMessageId; + if ( + parentMessageId != null && + removed.has(parentMessageId) && + !removed.has(message.messageId) + ) { + removed.add(message.messageId); + grew = true; + } + } } + return messages.filter((message) => !removed.has(message.messageId)); } return messages.filter((msg) => msg.messageId !== initialResponseId); @@ -192,6 +213,7 @@ export default function useChatFunctions({ const { getExpiry } = useUserKey(immutableConversation?.endpoint ?? ''); const setIsSubmitting = useSetRecoilState(store.isSubmittingFamily(index)); const setShowStopButton = useSetRecoilState(store.showStopButtonByIndex(index)); + const focusRegeneratedResponse = useFocusRegeneratedResponse(); /** * Atomically read + reset the per-conversation queue of manually-invoked @@ -554,6 +576,7 @@ export default function useChatFunctions({ if (isRegenerate) { setMessages([...submissionMessages, initialResponse]); + focusRegeneratedResponse(initialResponse.parentMessageId); } else { setMessages([...submissionMessages, currentMsg, initialResponse]); } diff --git a/client/src/hooks/Chat/useFocusRegeneratedResponse.ts b/client/src/hooks/Chat/useFocusRegeneratedResponse.ts new file mode 100644 index 0000000000..5d5ad3abc0 --- /dev/null +++ b/client/src/hooks/Chat/useFocusRegeneratedResponse.ts @@ -0,0 +1,25 @@ +import { useRecoilCallback } from 'recoil'; +import store from '~/store'; + +/** + * Focus a regenerated response's fork on its newest sibling. A regeneration + * appends the new response as the newest child of its parent, but when the + * parent already had multiple siblings the child count is unchanged (the slice + * keeps unrelated branches), so MultiMessage's length-change reset never fires + * and the view would otherwise stay on the kept sibling. Selecting reversed + * index 0 (newest = the appended response) makes the regenerating branch + * visible — applied at the optimistic render and reaffirmed on the `created` + * event so there is no flash on the kept sibling in between. + */ +export default function useFocusRegeneratedResponse() { + return useRecoilCallback( + ({ set }) => + (parentMessageId?: string | null) => { + if (parentMessageId == null) { + return; + } + set(store.messagesSiblingIdxFamily(parentMessageId), 0); + }, + [], + ); +} diff --git a/client/src/hooks/SSE/useEventHandlers.ts b/client/src/hooks/SSE/useEventHandlers.ts index 5cf7b5a41e..a01db036c9 100644 --- a/client/src/hooks/SSE/useEventHandlers.ts +++ b/client/src/hooks/SSE/useEventHandlers.ts @@ -38,6 +38,7 @@ import { queueTitleGeneration, markTitleGenerationProcessed, } from '~/data-provider'; +import useFocusRegeneratedResponse from '~/hooks/Chat/useFocusRegeneratedResponse'; import { shouldResetSubagentAtomsOnConversationChange } from './cleanup'; import useAttachmentHandler from '~/hooks/SSE/useAttachmentHandler'; import useContentHandler from '~/hooks/SSE/useContentHandler'; @@ -488,6 +489,8 @@ export default function useEventHandlers({ [queryClient, setMessages, isAddedRequest, announcePolite, setConversation, setShowStopButton], ); + const focusRegeneratedResponse = useFocusRegeneratedResponse(); + const createdHandler = useCallback( (data: TResData, submission: EventSubmission) => { queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]); @@ -510,6 +513,7 @@ export default function useEventHandlers({ }); if (isRegenerate) { setMessages([...messages, initialResponse]); + focusRegeneratedResponse(initialResponse.parentMessageId); } else { setMessages([...messages, userMessage, initialResponse]); } @@ -580,6 +584,7 @@ export default function useEventHandlers({ announcePolite, setConversation, applyAgentTemplate, + focusRegeneratedResponse, ], ); diff --git a/e2e/setup/fake-model.js b/e2e/setup/fake-model.js index 0e690a3679..77f1d98a34 100644 --- a/e2e/setup/fake-model.js +++ b/e2e/setup/fake-model.js @@ -23,6 +23,7 @@ const ASSERT_PROVIDER_FILE_MARKER = 'E2E_ASSERT_PROVIDER_FILE:'; const REPLY_MARKER = 'E2E_REPLY:'; const COUNTED_REPLY_MARKER = 'E2E_COUNTED_REPLY:'; const SLOW_REPLY_MARKER = 'E2E_SLOW_REPLY:'; +const SLOW_COUNTED_REPLY_MARKER = 'E2E_SLOW_COUNTED_REPLY:'; const RESUME_ICON_REPLY_MARKER = 'E2E_RESUME_ICON_REPLY:'; const FORCED_ERROR_MARKER = 'E2E_FORCED_ERROR:'; const MARKDOWN_REPLY_MARKER = 'E2E_MARKDOWN_REPLY'; @@ -49,6 +50,7 @@ const SKILL_DESCRIPTION = const EDITED_SKILL_DESCRIPTION = 'Use this edited skill to verify LibreChat skill file authoring in mock end-to-end tests.'; const countedReplies = new Map(); +const slowCountedReplies = new Map(); function messageType(message) { if (typeof message.getType === 'function') { @@ -285,6 +287,20 @@ function replyResponses(text) { }; } + const slowCountedName = getMarkerValue(text, SLOW_COUNTED_REPLY_MARKER); + if (slowCountedName) { + const count = (slowCountedReplies.get(slowCountedName) ?? 0) + 1; + slowCountedReplies.set(slowCountedName, count); + const chunks = Array.from( + { length: SLOW_REPLY_CHUNKS }, + (_, index) => `chunk-${String(index).padStart(3, '0')}`, + ).join(' '); + return { + responses: [`E2E slow counted reply ${slowCountedName} #${count} ${chunks}`], + sleep: SLOW_CHUNK_DELAY_MS, + }; + } + const resumeIconName = getMarkerValue(text, RESUME_ICON_REPLY_MARKER); if (resumeIconName) { const chunks = Array.from( diff --git a/e2e/specs/mock/chat.spec.ts b/e2e/specs/mock/chat.spec.ts index 0fcfb8bcbd..da72bdf161 100644 --- a/e2e/specs/mock/chat.spec.ts +++ b/e2e/specs/mock/chat.spec.ts @@ -6,6 +6,8 @@ import { NEW_CHAT_PATH, messagesView, mockReply, + replyText, + replyPrompt, selectMockEndpoint, sendMessage, } from './helpers'; @@ -188,6 +190,124 @@ test.describe('core chat loop', () => { await expect(page.getByText(followUpMessage)).toBeVisible(); }); + test('keeps the viewed branch when regenerating its latest response with an earlier branch present', async ({ + page, + }) => { + test.setTimeout(120000); + const firstMessage = 'first turn from e2e'; + const secondMessage = 'second turn from e2e'; + + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + + let response = await sendMessage(page, firstMessage); + expect(response.ok()).toBeTruthy(); + await expect(mockReply(page).first()).toBeVisible(); + response = await sendMessage(page, secondMessage); + expect(response.ok()).toBeTruthy(); + await expect(page.getByText(secondMessage)).toBeVisible(); + + // Regenerate the INITIAL response → a second root-level branch the second + // turn does not belong to. + const firstAssistant = messagesView(page).locator('.message-render').nth(1); + await firstAssistant.hover(); + const regenInitial = firstAssistant.locator('button[title="Regenerate"]').last(); + await expect(regenInitial).toBeVisible(); + [response] = await Promise.all([ + page.waitForResponse(isAgentsStream, { timeout: 30000 }), + regenInitial.click(), + ]); + expect(response.ok()).toBeTruthy(); + await expect(page.getByText('2 / 2')).toBeVisible(); + await expect(page.getByText(secondMessage)).toHaveCount(0); + + // Back to the ORIGINAL branch (both turns present). + await page.getByRole('button', { name: 'Previous sibling message' }).click(); + await expect(page.getByText('1 / 2')).toBeVisible(); + await expect(page.getByText(secondMessage)).toBeVisible(); + + // Regenerate the LATEST response on the original branch. The bug snapped the + // root fork back to the newest (regenerated-initial) branch, dropping the + // original thread; the view must stay put. + const latestAssistant = messagesView(page).locator('.message-render').last(); + await latestAssistant.hover(); + const regenLatest = latestAssistant.locator('button[title="Regenerate"]').last(); + await expect(regenLatest).toBeVisible(); + [response] = await Promise.all([ + page.waitForResponse(isAgentsStream, { timeout: 30000 }), + regenLatest.click(), + ]); + expect(response.ok()).toBeTruthy(); + + // Still on the original branch: the second turn survives and the root fork + // still reads 1 / 2 (rather than snapping to the regenerated-initial branch). + await expect(page.getByText(secondMessage)).toBeVisible(); + await expect(page.getByText('1 / 2')).toBeVisible(); + }); + + test('preserves a long original branch when regenerating early then later on it', async ({ + page, + }) => { + test.setTimeout(150000); + // Labeled prompts give each turn a unique reply, so we can both settle on + // it (turn complete) and assert which branch is visible. + const turns = [ + { prompt: replyPrompt('lb-one'), reply: replyText('lb-one') }, + { prompt: replyPrompt('lb-two'), reply: replyText('lb-two') }, + { prompt: replyPrompt('lb-three'), reply: replyText('lb-three') }, + ]; + + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + + // Build a three-turn thread (the "long running thread"), waiting for each + // turn's unique reply to render before sending the next. + for (const turn of turns) { + const response = await sendMessage(page, turn.prompt); + expect(response.ok()).toBeTruthy(); + await expect(messagesView(page).getByText(turn.reply)).toBeVisible({ timeout: 30000 }); + } + + // Regenerate from an EARLIER part of the branch (the first response). This + // forks a fresh root branch that does not contain the later turns. + const earlyAssistant = messagesView(page).locator('.message-render').nth(1); + await earlyAssistant.hover(); + const regenEarly = earlyAssistant.locator('button[title="Regenerate"]').last(); + await expect(regenEarly).toBeVisible(); + let [response] = await Promise.all([ + page.waitForResponse(isAgentsStream, { timeout: 30000 }), + regenEarly.click(), + ]); + expect(response.ok()).toBeTruthy(); + await expect(page.getByText('2 / 2')).toBeVisible(); + // The fresh branch does not contain the later turns' replies. + await expect(messagesView(page).getByText(turns[1].reply)).toHaveCount(0); + await expect(messagesView(page).getByText(turns[2].reply)).toHaveCount(0); + + // Go back to the ORIGINAL branch — all three turns are present again. + await page.getByRole('button', { name: 'Previous sibling message' }).click(); + await expect(page.getByText('1 / 2')).toBeVisible(); + await expect(messagesView(page).getByText(turns[1].reply)).toBeVisible(); + await expect(messagesView(page).getByText(turns[2].reply)).toBeVisible(); + + // Regenerate from LATER in the original branch (its latest response). The + // bug snapped the early fork back to the regenerated branch, collapsing the + // long original thread; it must stay intact. + const lateAssistant = messagesView(page).locator('.message-render').last(); + await lateAssistant.hover(); + const regenLate = lateAssistant.locator('button[title="Regenerate"]').last(); + await expect(regenLate).toBeVisible(); + [response] = await Promise.all([ + page.waitForResponse(isAgentsStream, { timeout: 30000 }), + regenLate.click(), + ]); + expect(response.ok()).toBeTruthy(); + + // The whole original branch survives and its first fork still reads 1 / 2. + await expect(messagesView(page).getByText(turns[1].reply)).toBeVisible(); + await expect(page.getByText('1 / 2')).toBeVisible(); + }); + test('keeps upload-to-provider CSV attached to the sent message and model input', async ({ page, }) => { diff --git a/e2e/specs/mock/message-tree.spec.ts b/e2e/specs/mock/message-tree.spec.ts index d9fd431ea5..f66c41dafa 100644 --- a/e2e/specs/mock/message-tree.spec.ts +++ b/e2e/specs/mock/message-tree.spec.ts @@ -51,6 +51,9 @@ const countedPrompt = (label: string) => `E2E_COUNTED_REPLY:${label}`; const countedReplyText = (label: string, count: number) => `E2E counted reply ${label} #${count}`; const slowPrompt = (label: string) => `E2E_SLOW_REPLY:${label}`; const slowReplyPrefix = (label: string) => `E2E slow reply ${label}`; +const slowCountedPrompt = (label: string) => `E2E_SLOW_COUNTED_REPLY:${label}`; +const slowCountedReplyText = (label: string, count: number) => + `E2E slow counted reply ${label} #${count}`; const messagesView = (page: Page) => page.getByTestId('messages-view'); const messageRender = (page: Page, text: string) => @@ -772,6 +775,104 @@ test.describe('message tree stream operations', () => { expect(messages.some((message) => messageText(message).includes(followReply))).toBe(false); }); + test('shows the regenerating response immediately when regenerating a non-latest sibling', async ({ + page, + }) => { + // A parent with multiple siblings, switched to an older one, then + // regenerated: the optimistic slice drops the target but keeps the other + // siblings, so the child count is unchanged and MultiMessage's + // length-change reset never fires. Without an explicit focus the view stays + // on the kept (newer) sibling and the streaming response is hidden until the + // server restores the dropped sibling at finalize. A slow reply keeps the + // stream open long enough to assert the during-stream state. + const label = uniqueLabel('regen-nonlatest'); + const prompt = slowCountedPrompt(label); + const reply1 = slowCountedReplyText(label, 1); + const reply2 = slowCountedReplyText(label, 2); + const reply3 = slowCountedReplyText(label, 3); + const stopButton = page.getByRole('button', { name: 'Stop generating' }); + + await openMockChat(page); + await sendMessage(page, prompt); + await expect(messagesView(page).getByText(reply1)).toBeVisible({ timeout: 30000 }); + await expect(stopButton).toBeHidden({ timeout: 30000 }); + + // First regenerate adds a sibling; the parent now has two responses. + await waitForGenerationStart(page, () => clickMessageTitleButton(page, reply1, 'Regenerate')); + await expect(messagesView(page).getByText(reply2)).toBeVisible({ timeout: 30000 }); + await expect(stopButton).toBeHidden({ timeout: 30000 }); + + // View the older sibling, then regenerate it (the non-latest one). + await clickSibling(page, reply2, 'Previous'); + await expect(messagesView(page).getByText(reply1)).toBeVisible(); + await expect(messagesView(page).getByText(reply2)).toBeHidden(); + + await waitForGenerationStart(page, () => clickMessageTitleButton(page, reply1, 'Regenerate')); + // Still streaming (slow reply): the regenerating response must be visible + // now — not the kept sibling — and well before finalize. + await expect(stopButton).toBeVisible({ timeout: 30000 }); + await expect(messagesView(page).getByText(reply3)).toBeVisible({ timeout: 4000 }); + await expect(messagesView(page).getByText(reply2)).toBeHidden(); + + // It stays put after finalize too. + await expect(stopButton).toBeHidden({ timeout: 30000 }); + await expect(messagesView(page).getByText(reply3)).toBeVisible(); + }); + + test('does not flash the kept sibling before the created event when regenerating a non-latest sibling', async ({ + page, + }) => { + // Same hazard at the optimistic (pre-`created`) render: useChatFunctions + // appends the placeholder and renders it before the request resolves. Gate + // the regenerate request so only that optimistic frame is on screen — with + // no `created` event to fall back on — and assert the kept sibling is + // already gone (the regenerating placeholder took its place). This isolates + // the optimistic focus; the createdHandler focus cannot mask a regression + // because `created` never arrives during the assertion. + const label = uniqueLabel('regen-precreated'); + const prompt = countedPrompt(label); + const reply1 = countedReplyText(label, 1); + const reply2 = countedReplyText(label, 2); + const reply3 = countedReplyText(label, 3); + + await openMockChat(page); + await sendAndExpectReply(page, prompt, reply1); + + await waitForGenerationStart(page, () => clickMessageTitleButton(page, reply1, 'Regenerate')); + await expect(messagesView(page).getByText(reply2)).toBeVisible({ timeout: 30000 }); + + await clickSibling(page, reply2, 'Previous'); + await expect(messagesView(page).getByText(reply1)).toBeVisible(); + await expect(messagesView(page).getByText(reply2)).toBeHidden(); + + // Hold the SSE stream so the `created` event cannot arrive: the optimistic + // render is all there is. The chat POST returns a stream id; the stream + // itself is a separate GET (/api/agents/chat/stream/) — gate that one, + // not the POST, which must complete to hand back the id. + let release = () => {}; + const gate = new Promise((resolve) => { + release = resolve; + }); + await page.route(/\/api\/agents\/chat\/stream\//, async (route: Route) => { + await gate; + return route.continue(); + }); + + try { + await clickMessageTitleButton(page, reply1, 'Regenerate'); + // The stream is gated, so the regenerating response has no text yet. + await expect(messagesView(page).getByText(reply3)).toBeHidden(); + // Pre-`created` window: the kept sibling must not be the one on screen — + // the regenerating placeholder has already taken its place. + await expect(messagesView(page).getByText(reply2)).toBeHidden({ timeout: 5000 }); + await expect(messagesView(page).getByText(reply1)).toBeHidden(); + } finally { + release(); + } + + await expect(messagesView(page).getByText(reply3)).toBeVisible({ timeout: 30000 }); + }); + test('resumes pending OAuth on the selected older branch after reload', async ({ page }) => { const label = uniqueLabel('oauth-branch'); const rootPrompt = countedPrompt(`${label}-root`);