diff --git a/client/src/hooks/Chat/useChatFunctions.ts b/client/src/hooks/Chat/useChatFunctions.ts index 1d4431940d..d65e38b522 100644 --- a/client/src/hooks/Chat/useChatFunctions.ts +++ b/client/src/hooks/Chat/useChatFunctions.ts @@ -26,9 +26,14 @@ import type { } from 'librechat-data-provider'; import type { SetterOrUpdater } from 'recoil'; import type { TAskFunction, ExtendedFile } from '~/common'; +import { + logger, + hasStreamStartFailed, + createDualMessageContent, + getRouteChatProjectId, +} from '~/utils'; import useSetFilesToDelete from '~/hooks/Files/useSetFilesToDelete'; import useGetSender from '~/hooks/Conversations/useGetSender'; -import { logger, createDualMessageContent, getRouteChatProjectId } from '~/utils'; import store, { useGetEphemeralAgent } from '~/store'; import { startupConfigKey } from '~/data-provider'; import useUserKey from '~/hooks/Input/useUserKey'; @@ -40,6 +45,31 @@ const logChatRequest = (request: Record) => { logger.log('====================================='); }; +const getAppendParentMessageId = ({ + latestMessage, + currentMessages, +}: { + latestMessage: TMessage | null; + currentMessages: TMessage[]; +}) => { + if (!latestMessage) { + return Constants.NO_PARENT; + } + + if (!hasStreamStartFailed(latestMessage)) { + return latestMessage.messageId; + } + + const failedUserMessage = currentMessages.find( + (message) => message.messageId === latestMessage.parentMessageId, + ); + if (failedUserMessage?.isCreatedByUser !== true) { + return latestMessage.messageId; + } + + return failedUserMessage.parentMessageId ?? Constants.NO_PARENT; +}; + export default function useChatFunctions({ index = 0, files, @@ -165,7 +195,7 @@ export default function useChatFunctions({ } const isEditOrContinue = isEdited || isContinued; - let currentMessages: TMessage[] | null = overrideMessages ?? getMessages() ?? []; + let currentMessages: TMessage[] = overrideMessages ?? getMessages() ?? []; if (conversation?.promptPrefix) { conversation.promptPrefix = replaceSpecialVars({ @@ -184,7 +214,8 @@ export default function useChatFunctions({ // construct the query message // this is not a real messageId, it is used as placeholder before real messageId returned const intermediateId = overrideUserMessageId ?? v4(); - parentMessageId = parentMessageId ?? latestMessage?.messageId ?? Constants.NO_PARENT; + parentMessageId = + parentMessageId ?? getAppendParentMessageId({ latestMessage, currentMessages }); logChatRequest({ index, diff --git a/client/src/hooks/Input/useTextarea.ts b/client/src/hooks/Input/useTextarea.ts index fe65208ba0..b8cd22cfbc 100644 --- a/client/src/hooks/Input/useTextarea.ts +++ b/client/src/hooks/Input/useTextarea.ts @@ -1,5 +1,5 @@ -import debounce from 'lodash/debounce'; import { useEffect, useRef, useCallback } from 'react'; +import debounce from 'lodash/debounce'; import { useRecoilValue, useRecoilState } from 'recoil'; import type { TEndpointOption } from 'librechat-data-provider'; import type { KeyboardEvent } from 'react'; @@ -11,12 +11,12 @@ import { checkIfScrollable, } from '~/utils'; import { useAssistantsMapContext } from '~/Providers/AssistantsMapContext'; +import { useLatestMessage } from '~/hooks/Messages/useLatestMessage'; import { useAgentsMapContext } from '~/Providers/AgentsMapContext'; import useGetSender from '~/hooks/Conversations/useGetSender'; import useFileHandling from '~/hooks/Files/useFileHandling'; import { useInteractionHealthCheck } from '~/data-provider'; import { useChatContext } from '~/Providers/ChatContext'; -import { useLatestMessage } from '~/hooks/Messages/useLatestMessage'; import { globalAudioId } from '~/common'; import { useLocalize } from '~/hooks'; import store from '~/store'; @@ -59,7 +59,8 @@ export default function useTextarea({ }); const entityName = entity?.name ?? ''; - const isNotAppendable = latestMessage?.error === true && !isAssistant; + const isNotAppendable = + latestMessage?.error === true && latestMessage.isCreatedByUser === true && !isAssistant; // && (conversationId?.length ?? 0) > 6; // also ensures that we don't show the wrong placeholder useEffect(() => { diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts index c1b8e3fc6c..fec176b456 100644 --- a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts +++ b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts @@ -1,4 +1,4 @@ -import { renderHook, act } from '@testing-library/react'; +import { renderHook, act, waitFor } from '@testing-library/react'; import { Constants, LocalStorageKeys, QueryKeys, request } from 'librechat-data-provider'; import type { TSubmission } from 'librechat-data-provider'; @@ -48,6 +48,33 @@ const mockQueryClient = { }), }; +const mockActiveRunAtom = { key: 'activeRun' }; +const mockAbortScrollAtom = { key: 'abortScroll' }; +const mockSubmissionAtom = { key: 'submission' }; +const mockShowStopButtonAtom = { key: 'showStopButton' }; +const mockSetActiveRun = jest.fn(); +const mockSetAbortScroll = jest.fn(); +const mockSetSubmission = jest.fn(); +const mockSetShowStopButton = jest.fn(); +const mockUseSetRecoilStateMock = jest.fn((atom: unknown) => { + if (atom === mockActiveRunAtom) { + return mockSetActiveRun; + } + if (atom === mockAbortScrollAtom) { + return mockSetAbortScroll; + } + if (atom === mockSubmissionAtom) { + return mockSetSubmission; + } + if (atom === mockShowStopButtonAtom) { + return mockSetShowStopButton; + } + return jest.fn(); +}); +function mockUseSetRecoilState(atom: unknown) { + return mockUseSetRecoilStateMock(atom); +} + jest.mock('@tanstack/react-query', () => ({ ...jest.requireActual('@tanstack/react-query'), useQueryClient: () => mockQueryClient, @@ -55,15 +82,16 @@ jest.mock('@tanstack/react-query', () => ({ jest.mock('recoil', () => ({ ...jest.requireActual('recoil'), - useSetRecoilState: () => jest.fn(), + useSetRecoilState: mockUseSetRecoilState, })); jest.mock('~/store', () => ({ __esModule: true, default: { - activeRunFamily: jest.fn(), - abortScrollFamily: jest.fn(), - showStopButtonByIndex: jest.fn(), + activeRunFamily: jest.fn(() => mockActiveRunAtom), + abortScrollFamily: jest.fn(() => mockAbortScrollAtom), + submissionByIndex: jest.fn(() => mockSubmissionAtom), + showStopButtonByIndex: jest.fn(() => mockShowStopButtonAtom), }, })); @@ -211,6 +239,11 @@ describe('useResumableSSE - 404 error path', () => { mockInvalidateQueries.mockClear(); mockRemoveQueries.mockClear(); mockFindAll.mockClear(); + mockUseSetRecoilStateMock.mockClear(); + mockSetActiveRun.mockClear(); + mockSetAbortScroll.mockClear(); + mockSetSubmission.mockClear(); + mockSetShowStopButton.mockClear(); (request.post as jest.Mock).mockReset(); (request.post as jest.Mock).mockResolvedValue({ streamId: 'stream-123' }); }); @@ -537,6 +570,39 @@ describe('useResumableSSE - 404 error path', () => { jest.useRealTimers(); }); + it('clears submission and stop state when starting generation fails', async () => { + (request.post as jest.Mock).mockRejectedValueOnce({ + response: { + status: 500, + data: { message: 'failed to start' }, + }, + }); + const submission = buildSubmission(); + const chatHelpers = buildChatHelpers(); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + + await waitFor(() => { + expect(mockSetSubmission).toHaveBeenCalledWith(null); + }); + + expect(mockSSEInstances).toHaveLength(0); + expect(mockErrorHandler).toHaveBeenCalledWith( + expect.objectContaining({ + data: { + text: JSON.stringify({ message: 'failed to start' }), + metadata: { streamStartFailed: true }, + }, + submission, + }), + ); + expect(mockSetIsSubmitting).toHaveBeenCalledWith(true); + expect(mockSetIsSubmitting).toHaveBeenCalledWith(false); + expect(mockSetShowStopButton).toHaveBeenCalledWith(true); + expect(mockSetShowStopButton).toHaveBeenCalledWith(false); + unmount(); + }); + it('replays title events from resume state sync', async () => { const submission = buildSubmission(); const chatHelpers = buildChatHelpers(); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 0bbbca8878..198b1b2c0d 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -22,16 +22,22 @@ import type { EventSubmission, } from 'librechat-data-provider'; import type { EventHandlerParams } from './useEventHandlers'; +import type { ActiveJobsResponse } from '~/data-provider'; +import type { TResData } from '~/common'; +import { + clearAllDrafts, + removeConvoFromAllQueries, + upsertConvoInAllQueries, + markStreamStartFailedMetadata, +} from '~/utils'; import { useGetUserBalance, useGetStartupConfig, queueTitleGeneration, streamStatusQueryKey, } from '~/data-provider'; -import type { ActiveJobsResponse } from '~/data-provider'; import { useAuthContext } from '~/hooks/AuthContext'; import useEventHandlers from './useEventHandlers'; -import { clearAllDrafts, removeConvoFromAllQueries, upsertConvoInAllQueries } from '~/utils'; import store from '~/store'; type ChatHelpers = Pick< @@ -39,6 +45,14 @@ type ChatHelpers = Pick< 'setMessages' | 'getMessages' | 'setConversation' | 'setIsSubmitting' | 'newConversation' >; +const getStreamStartFailureData = (errorData?: Record): TResData => + ({ + text: errorData + ? JSON.stringify(errorData) + : 'Error connecting to server, try refreshing the page.', + metadata: markStreamStartFailedMetadata(), + }) as unknown as TResData; + const MAX_RETRIES = 5; const START_GENERATION_NETWORK_RETRIES = 3; const START_GENERATION_READINESS_TIMEOUT_MS = 120000; @@ -263,6 +277,7 @@ export default function useResumableSSE( const [_completed, setCompleted] = useState(new Set()); const [streamId, setStreamId] = useState(null); const setAbortScroll = useSetRecoilState(store.abortScrollFamily(runIndex)); + const setSubmission = useSetRecoilState(store.submissionByIndex(runIndex)); const setShowStopButton = useSetRecoilState(store.showStopButtonByIndex(runIndex)); const sseRef = useRef(null); @@ -849,20 +864,16 @@ export default function useResumableSSE( const axiosError = lastError as { response?: { data?: Record } }; const errorData = axiosError?.response?.data; - if (errorData) { - errorHandler({ - data: { text: JSON.stringify(errorData) } as unknown as Parameters< - typeof errorHandler - >[0]['data'], - submission: currentSubmission as EventSubmission, - }); - } else { - errorHandler({ data: undefined, submission: currentSubmission as EventSubmission }); - } + errorHandler({ + data: getStreamStartFailureData(errorData), + submission: currentSubmission as EventSubmission, + }); + setShowStopButton(false); setIsSubmitting(false); + setSubmission(null); return null; }, - [clearStepMaps, errorHandler, setIsSubmitting], + [clearStepMaps, errorHandler, setIsSubmitting, setShowStopButton, setSubmission], ); useEffect(() => { diff --git a/client/src/utils/messages.ts b/client/src/utils/messages.ts index d906c33ddc..40507f3b16 100644 --- a/client/src/utils/messages.ts +++ b/client/src/utils/messages.ts @@ -17,6 +17,7 @@ import type { QueryClient } from '@tanstack/react-query'; import type { LocalizeFunction } from '~/common'; export const TEXT_KEY_DIVIDER = '|||'; +export const STREAM_START_FAILED_METADATA_KEY = 'streamStartFailed'; type SiblingIndexLookup = (parentMessageId: string | null | undefined) => number; @@ -136,6 +137,16 @@ export const getAllContentText = (message?: TMessage | null): string => { return ''; }; +export const hasStreamStartFailed = (message?: Pick | null): boolean => + message?.metadata?.[STREAM_START_FAILED_METADATA_KEY] === true; + +export const markStreamStartFailedMetadata = ( + metadata?: TMessage['metadata'], +): TMessage['metadata'] => ({ + ...(metadata ?? {}), + [STREAM_START_FAILED_METADATA_KEY]: true, +}); + const getLatestContentForKey = (message: TMessage): string => { const formatText = (str: string, index: number): string => { if (str.length === 0) { diff --git a/e2e/setup/fake-model.js b/e2e/setup/fake-model.js index eea805be5d..5b6a58191c 100644 --- a/e2e/setup/fake-model.js +++ b/e2e/setup/fake-model.js @@ -9,6 +9,8 @@ * without a live provider or a standalone HTTP mock server: responses are decided * from the conversation and the agents' advertised tools. */ +const { FakeChatModel } = require('@librechat/agents'); + const MOCK_REPLY = process.env.MOCK_LLM_REPLY || 'E2E mock reply: pong'; const CHUNK_DELAY_MS = Number(process.env.MOCK_LLM_CHUNK_DELAY_MS) || 10; @@ -16,10 +18,16 @@ const CREATE_SKILL_MARKER = 'E2E_CREATE_SKILL:'; const EDIT_SKILL_MARKER = 'E2E_EDIT_SKILL:'; const ASSERT_MODEL_SPEC_SKILLS_MARKER = 'E2E_ASSERT_MODEL_SPEC_SKILLS'; 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 FORCED_ERROR_MARKER = 'E2E_FORCED_ERROR:'; const CREATE_FILE_AUTHORING_FINAL_TEXT = 'E2E file authoring complete'; const EDIT_FILE_AUTHORING_FINAL_TEXT = 'E2E file edit complete'; const MODEL_SPEC_SKILL_ASSERTION_FINAL_TEXT = 'E2E model spec skill assertion passed'; const PROVIDER_FILE_ASSERTION_FINAL_TEXT = 'E2E provider file assertion passed'; +const SLOW_CHUNK_DELAY_MS = Number(process.env.MOCK_LLM_SLOW_CHUNK_DELAY_MS) || 35; +const SLOW_REPLY_CHUNKS = 160; const CREATE_FILE_TOOL_NAME = 'create_file'; const EDIT_FILE_TOOL_NAME = 'edit_file'; const BASH_TOOL_NAME = 'bash_tool'; @@ -34,6 +42,7 @@ const SKILL_DESCRIPTION = 'Use this skill to verify LibreChat skill file authoring in mock end-to-end tests.'; const EDITED_SKILL_DESCRIPTION = 'Use this edited skill to verify LibreChat skill file authoring in mock end-to-end tests.'; +const countedReplies = new Map(); function messageType(message) { if (typeof message.getType === 'function') { @@ -101,7 +110,12 @@ function getMarkerValue(text, marker) { if (markerIndex === -1) { return ''; } - return text.slice(markerIndex + marker.length).trim().split(/\s+/, 1)[0] ?? ''; + return ( + text + .slice(markerIndex + marker.length) + .trim() + .split(/\s+/, 1)[0] ?? '' + ); } function collectToolNames(agents) { @@ -210,6 +224,70 @@ function providerFileAssertionResponses({ messages, text }) { }; } +function replyResponses(text) { + const errorName = getMarkerValue(text, FORCED_ERROR_MARKER); + if (errorName) { + return { + responses: [`E2E forced error prelude ${errorName}`], + thrownError: `E2E forced stream error ${errorName}`, + }; + } + + const replyName = getMarkerValue(text, REPLY_MARKER); + if (replyName) { + return { + responses: [`E2E reply ${replyName}`], + }; + } + + const countedName = getMarkerValue(text, COUNTED_REPLY_MARKER); + if (countedName) { + const count = (countedReplies.get(countedName) ?? 0) + 1; + countedReplies.set(countedName, count); + return { + responses: [`E2E counted reply ${countedName} #${count}`], + }; + } + + const slowName = getMarkerValue(text, SLOW_REPLY_MARKER); + if (slowName) { + const chunks = Array.from( + { length: SLOW_REPLY_CHUNKS }, + (_, index) => `chunk-${String(index).padStart(3, '0')}`, + ).join(' '); + return { + responses: [`E2E slow reply ${slowName} ${chunks}`], + sleep: SLOW_CHUNK_DELAY_MS, + }; + } + + return null; +} + +function overrideModel({ graph, responses, sleep, toolCalls, thrownError }) { + if (!thrownError) { + graph.overrideTestModel(responses, sleep ?? CHUNK_DELAY_MS, toolCalls); + return; + } + + class ThrowingFakeChatModel extends FakeChatModel { + async *_streamResponseChunks(messages, options, runManager) { + yield* super._streamResponseChunks( + messages, + { ...options, thrownErrorString: thrownError }, + runManager, + ); + } + } + + graph.overrideModel = new ThrowingFakeChatModel({ + responses, + sleep: sleep ?? CHUNK_DELAY_MS, + emitCustomEvent: true, + toolCalls, + }); +} + function modelSpecSkillAssertionResponses({ agents, messages, toolNames }) { const failures = []; const additionalInstructions = collectAdditionalInstructions(agents); @@ -311,6 +389,11 @@ function fileAuthoringResponses(operation, toolNames) { } function resolveResponses({ agents, messages, text, toolNames }) { + const reply = replyResponses(text); + if (reply) { + return reply; + } + const providerFileAssertion = providerFileAssertionResponses({ messages, text }); if (providerFileAssertion) { return providerFileAssertion; @@ -361,11 +444,11 @@ module.exports = function fakeModelHook(run, context) { const text = getLatestUserText(context?.messages); const toolNames = collectToolNames(context?.agents); - const { responses, toolCalls } = resolveResponses({ + const { responses, sleep, toolCalls, thrownError } = resolveResponses({ agents: context?.agents, messages: context?.messages, text, toolNames, }); - graph.overrideTestModel(responses, CHUNK_DELAY_MS, toolCalls); + overrideModel({ graph, responses, sleep, toolCalls, thrownError }); }; diff --git a/e2e/specs/mock/chat.spec.ts b/e2e/specs/mock/chat.spec.ts index 9d08c845a8..f189003484 100644 --- a/e2e/specs/mock/chat.spec.ts +++ b/e2e/specs/mock/chat.spec.ts @@ -87,6 +87,7 @@ test.describe('core chat loop', () => { ).toBeVisible(); await expect(fileChip).toBeVisible(); + await expect(page).toHaveURL(/\/c\/[0-9a-fA-F-]{36}$/); const conversationUrl = page.url(); await page.reload({ timeout: 10000 }); await expect(page).toHaveURL(conversationUrl); diff --git a/e2e/specs/mock/helpers.ts b/e2e/specs/mock/helpers.ts index 0ec03b16b0..2ccae83ce2 100644 --- a/e2e/specs/mock/helpers.ts +++ b/e2e/specs/mock/helpers.ts @@ -19,7 +19,18 @@ type RefreshTokenBody = { }; export function isAgentsStream(response: Response) { - return response.url().includes('/api/agents') && response.status() === 200; + return isAgentGenerationStart(response); +} + +export function isAgentGenerationStart(response: Response) { + const { pathname } = new URL(response.url()); + const isAgentsChat = pathname === '/api/agents/chat' || pathname.startsWith('/api/agents/chat/'); + return ( + response.request().method() === 'POST' && + isAgentsChat && + !pathname.endsWith('/abort') && + response.status() === 200 + ); } const modelSelectorTrigger = (page: Page) => diff --git a/e2e/specs/mock/message-tree.spec.ts b/e2e/specs/mock/message-tree.spec.ts new file mode 100644 index 0000000000..bb7096dc36 --- /dev/null +++ b/e2e/specs/mock/message-tree.spec.ts @@ -0,0 +1,635 @@ +import { expect, test } from '@playwright/test'; +import type { Page, Response, Route } from '@playwright/test'; +import { + isAgentGenerationStart, + MOCK_ENDPOINTS, + NEW_CHAT_PATH, + fetchJson, + getAccessToken, + selectMockEndpoint, + sendMessage, +} from './helpers'; + +const NO_PARENT = '00000000-0000-0000-0000-000000000000'; + +type TextContentPart = { + type?: string; + text?: string | { value?: string }; + error?: string; +}; + +type E2EMessage = { + messageId: string; + parentMessageId?: string | null; + conversationId?: string | null; + text?: string; + content?: TextContentPart[]; + isCreatedByUser?: boolean; + error?: boolean; + unfinished?: boolean; +}; + +type ForkResponse = { + conversation: { + conversationId?: string; + }; + messages: E2EMessage[]; +}; + +const uniqueLabel = (name: string) => `${name}-${Date.now()}-${Math.floor(Math.random() * 1e6)}`; + +const replyPrompt = (label: string) => `E2E_REPLY:${label}`; +const replyText = (label: string) => `E2E reply ${label}`; +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 messagesView = (page: Page) => page.getByTestId('messages-view'); +const messageRender = (page: Page, text: string) => + page.locator('.message-render').filter({ hasText: text }).last(); +const conversationPath = (conversationId: string) => `/c/${encodeURIComponent(conversationId)}`; + +function contentText(part: TextContentPart): string { + if (typeof part.text === 'string') { + return part.text; + } + if (part.text?.value) { + return part.text.value; + } + return part.error ?? ''; +} + +function messageText(message: E2EMessage): string { + if (message.text) { + return message.text; + } + return message.content?.map(contentText).filter(Boolean).join('\n') ?? ''; +} + +function findMessage(messages: E2EMessage[], text: string, isCreatedByUser?: boolean): E2EMessage { + const message = messages.find((candidate) => { + const roleMatches = + isCreatedByUser === undefined || candidate.isCreatedByUser === isCreatedByUser; + return roleMatches && messageText(candidate).includes(text); + }); + if (!message) { + throw new Error( + `Expected message containing "${text}". Saw:\n${messages.map(messageText).join('\n---\n')}`, + ); + } + return message; +} + +function expectParent( + messages: E2EMessage[], + childText: string, + parentText: string, + childIsUser?: boolean, +) { + const child = findMessage(messages, childText, childIsUser); + const parent = findMessage(messages, parentText); + expect(child.parentMessageId, `${childText} should be a child of ${parentText}`).toBe( + parent.messageId, + ); +} + +function expectNoFoldedMessages(messages: E2EMessage[]) { + const ids = new Set(messages.map((message) => message.messageId)); + const folded = messages.filter((message) => { + const parentId = message.parentMessageId; + return parentId != null && parentId !== '' && parentId !== NO_PARENT && !ids.has(parentId); + }); + expect( + folded.map((message) => ({ + text: messageText(message), + messageId: message.messageId, + parentMessageId: message.parentMessageId, + })), + 'messages must not render as parent-less folded children', + ).toEqual([]); + + const roots = messages.filter((message) => { + const parentId = message.parentMessageId; + return parentId == null || parentId === '' || parentId === NO_PARENT; + }); + expect( + roots.map((message) => ({ + text: messageText(message), + isCreatedByUser: message.isCreatedByUser, + })), + 'only user messages should be roots', + ).toEqual(roots.map(() => expect.objectContaining({ isCreatedByUser: true }))); +} + +async function expectVisibleMessages(page: Page, texts: string[]) { + for (const text of texts) { + await expect(messagesView(page).getByText(text)).toBeVisible({ timeout: 30000 }); + } +} + +async function reloadAndExpectMessages(page: Page, texts: string[]) { + await page.reload({ timeout: 10000 }); + await expectVisibleMessages(page, texts); +} + +async function revisitConversationAndExpectMessages( + page: Page, + conversationId: string, + texts: string[], +) { + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await page.goto(conversationPath(conversationId), { timeout: 10000 }); + await expectVisibleMessages(page, texts); +} + +async function openMockChat(page: Page) { + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); +} + +function isAgentGenerationResponse(response: Response, expectedStatus: number) { + const { pathname } = new URL(response.url()); + const isAgentsChat = pathname === '/api/agents/chat' || pathname.startsWith('/api/agents/chat/'); + return ( + response.request().method() === 'POST' && + isAgentsChat && + !pathname.endsWith('/abort') && + response.status() === expectedStatus + ); +} + +async function waitForGenerationStart(page: Page, action: () => Promise): Promise { + const [response] = await Promise.all([ + page.waitForResponse(isAgentGenerationStart, { timeout: 30000 }), + action(), + ]); + expect(response.ok()).toBeTruthy(); + return response; +} + +async function sendAndExpectReply(page: Page, prompt: string, expectedReply: string) { + const response = await sendMessage(page, prompt); + expect(response.ok()).toBeTruthy(); + await expect(messagesView(page).getByText(expectedReply)).toBeVisible({ timeout: 30000 }); +} + +async function submitMessageExpectingGenerationFailure( + page: Page, + prompt: string, + expectedStatus: number, +) { + const input = page.getByRole('textbox', { name: 'Message input' }); + await expect(input).toBeEnabled({ timeout: 30000 }); + await input.click(); + await input.fill(prompt); + const [response] = await Promise.all([ + page.waitForResponse((res) => isAgentGenerationResponse(res, expectedStatus), { + timeout: 30000, + }), + input.press('Enter'), + ]); + return response; +} + +async function conversationIdFromPage(page: Page): Promise { + await expect(page).toHaveURL(/\/c\/(?!new)[0-9a-fA-F-]{36}$/); + const id = new URL(page.url()).pathname.split('/').pop(); + if (!id) { + throw new Error(`Could not parse conversation id from ${page.url()}`); + } + return id; +} + +async function fetchMessages( + page: Page, + conversationId: string, + accessToken?: string, +): Promise { + const token = accessToken ?? (await getAccessToken(page)); + return fetchJson( + page, + `/api/messages/${encodeURIComponent(conversationId)}`, + token, + ); +} + +async function waitForMessages( + page: Page, + conversationId: string, + predicate: (messages: E2EMessage[]) => boolean, + description: string, +): Promise { + let latest: E2EMessage[] = []; + const token = await getAccessToken(page); + for (let attempt = 0; attempt < 80; attempt++) { + latest = await fetchMessages(page, conversationId, token); + if (predicate(latest)) { + return latest; + } + await new Promise((resolve) => setTimeout(resolve, 250)); + } + + throw new Error( + `Timed out waiting for ${description}. Latest messages:\n${latest + .map( + (message) => `${message.messageId} <- ${message.parentMessageId}: ${messageText(message)}`, + ) + .join('\n')}`, + ); +} + +async function clickMessageTitleButton(page: Page, messageTextValue: string, title: string) { + const render = messageRender(page, messageTextValue); + await render.scrollIntoViewIfNeeded(); + await render.hover(); + await render.locator(`button[title="${title}"]`).last().click(); +} + +async function clickSibling(page: Page, messageTextValue: string, direction: 'Previous' | 'Next') { + const render = messageRender(page, messageTextValue); + await render.scrollIntoViewIfNeeded(); + await render.hover(); + await render.getByRole('button', { name: `${direction} sibling message` }).click(); +} + +async function expectCanCycleSiblingTexts(page: Page, previousText: string, nextText: string) { + const previous = messagesView(page).getByText(previousText); + const next = messagesView(page).getByText(nextText); + if (await previous.isVisible()) { + await clickSibling(page, previousText, 'Next'); + await expect(next).toBeVisible(); + await clickSibling(page, nextText, 'Previous'); + await expect(previous).toBeVisible(); + return; + } + + if (await next.isVisible()) { + await clickSibling(page, nextText, 'Previous'); + await expect(previous).toBeVisible(); + await clickSibling(page, previousText, 'Next'); + await expect(next).toBeVisible(); + return; + } + + throw new Error(`Expected either sibling "${previousText}" or "${nextText}" to be visible`); +} + +async function clickForkVisibleMessages( + page: Page, + messageTextValue: string, +): Promise { + const render = messageRender(page, messageTextValue); + await render.scrollIntoViewIfNeeded(); + await render.hover(); + await render.getByRole('button', { name: 'Open Fork Menu' }).click(); + + const [response] = await Promise.all([ + page.waitForResponse( + (res) => + res.request().method() === 'POST' && + res.url().includes('/api/convos/fork') && + res.status() === 200, + { timeout: 30000 }, + ), + page.getByRole('button', { name: 'Visible messages only', exact: true }).click(), + ]); + + return (await response.json()) as ForkResponse; +} + +test.describe('message tree stream operations', () => { + test.setTimeout(180000); + + test('streams follow-ups and keeps an aborted response as the next parent', async ({ page }) => { + const label = uniqueLabel('abort'); + const firstPrompt = replyPrompt(`${label}-first`); + const firstReply = replyText(`${label}-first`); + const secondPrompt = replyPrompt(`${label}-second`); + const secondReply = replyText(`${label}-second`); + const abortPrompt = slowPrompt(`${label}-stop`); + const abortReply = slowReplyPrefix(`${label}-stop`); + const afterAbortPrompt = replyPrompt(`${label}-after-stop`); + const afterAbortReply = replyText(`${label}-after-stop`); + + await openMockChat(page); + await sendAndExpectReply(page, firstPrompt, firstReply); + const conversationId = await conversationIdFromPage(page); + await sendAndExpectReply(page, secondPrompt, secondReply); + + const slowStart = await sendMessage(page, abortPrompt); + expect(slowStart.ok()).toBeTruthy(); + await expect(messagesView(page).getByText(abortReply)).toBeVisible({ timeout: 30000 }); + + const [abortResponse] = await Promise.all([ + page.waitForResponse( + (response) => + response.request().method() === 'POST' && + response.url().includes('/api/agents/chat/abort'), + { timeout: 30000 }, + ), + page.getByRole('button', { name: 'Stop generating' }).click(), + ]); + expect(abortResponse.ok()).toBeTruthy(); + await expect(page.getByRole('button', { name: 'Stop generating' })).toBeHidden({ + timeout: 30000, + }); + + let messages = await waitForMessages( + page, + conversationId, + (items) => items.some((message) => messageText(message).includes(abortReply)), + 'aborted response to persist', + ); + expectNoFoldedMessages(messages); + expectParent(messages, secondPrompt, firstReply, true); + expectParent(messages, abortReply, abortPrompt, false); + + await sendAndExpectReply(page, afterAbortPrompt, afterAbortReply); + messages = await waitForMessages( + page, + conversationId, + (items) => items.some((message) => messageText(message).includes(afterAbortReply)), + 'follow-up after abort', + ); + expectNoFoldedMessages(messages); + expectParent(messages, afterAbortPrompt, abortReply, true); + expectParent(messages, afterAbortReply, afterAbortPrompt, false); + + await reloadAndExpectMessages(page, [firstReply, secondReply, abortReply, afterAbortReply]); + await revisitConversationAndExpectMessages(page, conversationId, [ + firstReply, + secondReply, + abortReply, + afterAbortReply, + ]); + }); + + test('regenerates assistant siblings, cycles branches, follows up, and forks the visible branch', async ({ + page, + }) => { + const label = uniqueLabel('regen'); + const prompt = 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, prompt, firstReply); + const originalConversationId = await conversationIdFromPage(page); + + await waitForGenerationStart(page, () => + clickMessageTitleButton(page, firstReply, 'Regenerate'), + ); + await expect(messagesView(page).getByText(regeneratedReply)).toBeVisible({ timeout: 30000 }); + + await clickSibling(page, regeneratedReply, 'Previous'); + await expect(messagesView(page).getByText(firstReply)).toBeVisible(); + await expect(messagesView(page).getByText(regeneratedReply)).toBeHidden(); + await clickSibling(page, firstReply, 'Next'); + await expect(messagesView(page).getByText(regeneratedReply)).toBeVisible(); + + await sendAndExpectReply(page, followPrompt, followReply); + let messages = await waitForMessages( + page, + originalConversationId, + (items) => items.some((message) => messageText(message).includes(followReply)), + 'follow-up after regenerate', + ); + expectNoFoldedMessages(messages); + expectParent(messages, firstReply, prompt, false); + expectParent(messages, regeneratedReply, prompt, false); + expectParent(messages, followPrompt, regeneratedReply, true); + expectParent(messages, followReply, followPrompt, false); + + const userMessage = findMessage(messages, prompt, true); + const assistantSiblings = messages.filter( + (message) => message.parentMessageId === userMessage.messageId && !message.isCreatedByUser, + ); + expect(assistantSiblings.map(messageText).sort()).toEqual( + [firstReply, regeneratedReply].sort(), + ); + + await reloadAndExpectMessages(page, [regeneratedReply, followReply]); + await revisitConversationAndExpectMessages(page, originalConversationId, [ + regeneratedReply, + followReply, + ]); + await clickSibling(page, regeneratedReply, 'Previous'); + await expect(messagesView(page).getByText(firstReply)).toBeVisible(); + const fork = await clickForkVisibleMessages(page, firstReply); + const forkedConversationId = fork.conversation.conversationId; + if (!forkedConversationId) { + throw new Error('Expected fork response to include a conversation id'); + } + await expect(page).toHaveURL(new RegExp(`/c/${forkedConversationId}$`)); + + messages = fork.messages; + expectNoFoldedMessages(messages); + expect(messages.some((message) => messageText(message).includes(firstReply))).toBe(true); + expect(messages.some((message) => messageText(message).includes(regeneratedReply))).toBe(false); + expect(messages.some((message) => messageText(message).includes(followReply))).toBe(false); + }); + + test('long threads retain regenerated and save-and-submit branches after revisit', async ({ + page, + }) => { + const label = uniqueLabel('save-submit'); + const rootPrompt = replyPrompt(`${label}-root`); + const rootReply = replyText(`${label}-root`); + const firstPrompt = replyPrompt(`${label}-first`); + const firstReply = replyText(`${label}-first`); + const middlePrompt = replyPrompt(`${label}-middle`); + const middleReply = replyText(`${label}-middle`); + const fourthPrompt = replyPrompt(`${label}-fourth`); + const fourthReply = replyText(`${label}-fourth`); + const tailPrompt = countedPrompt(`${label}-tail`); + const tailReply = countedReplyText(`${label}-tail`, 1); + const regeneratedTailReply = countedReplyText(`${label}-tail`, 2); + const editedMiddlePrompt = replyPrompt(`${label}-middle-edited`); + const editedMiddleReply = replyText(`${label}-middle-edited`); + const afterEditPrompt = replyPrompt(`${label}-after-edit`); + const afterEditReply = replyText(`${label}-after-edit`); + + await openMockChat(page); + await sendAndExpectReply(page, rootPrompt, rootReply); + const conversationId = await conversationIdFromPage(page); + await sendAndExpectReply(page, firstPrompt, firstReply); + await sendAndExpectReply(page, middlePrompt, middleReply); + await sendAndExpectReply(page, fourthPrompt, fourthReply); + await sendAndExpectReply(page, tailPrompt, tailReply); + + await waitForGenerationStart(page, () => + clickMessageTitleButton(page, tailReply, 'Regenerate'), + ); + await expect(messagesView(page).getByText(regeneratedTailReply)).toBeVisible({ + timeout: 30000, + }); + + await clickMessageTitleButton(page, middlePrompt, 'Edit'); + const editor = page.getByTestId('message-text-editor'); + await expect(editor).toBeVisible(); + await editor.fill(editedMiddlePrompt); + await waitForGenerationStart(page, () => + page.getByRole('button', { name: 'Save & Submit' }).click(), + ); + await expect(messagesView(page).getByText(editedMiddleReply)).toBeVisible({ timeout: 30000 }); + + let messages = await waitForMessages( + page, + conversationId, + (items) => items.some((message) => messageText(message).includes(editedMiddleReply)), + 'save-and-submit edited branch', + ); + expectNoFoldedMessages(messages); + expectParent(messages, firstPrompt, rootReply, true); + expectParent(messages, firstReply, firstPrompt, false); + expectParent(messages, middlePrompt, firstReply, true); + expectParent(messages, middleReply, middlePrompt, false); + expectParent(messages, fourthPrompt, middleReply, true); + expectParent(messages, fourthReply, fourthPrompt, false); + expectParent(messages, tailPrompt, fourthReply, true); + expectParent(messages, tailReply, tailPrompt, false); + expectParent(messages, regeneratedTailReply, tailPrompt, false); + expectParent(messages, editedMiddlePrompt, firstReply, true); + expectParent(messages, editedMiddleReply, editedMiddlePrompt, false); + + await clickSibling(page, editedMiddlePrompt, 'Previous'); + await expectVisibleMessages(page, [middlePrompt, fourthReply]); + await expectCanCycleSiblingTexts(page, tailReply, regeneratedTailReply); + await clickSibling(page, middlePrompt, 'Next'); + await expectVisibleMessages(page, [editedMiddlePrompt, editedMiddleReply]); + + await sendAndExpectReply(page, afterEditPrompt, afterEditReply); + messages = await waitForMessages( + page, + conversationId, + (items) => items.some((message) => messageText(message).includes(afterEditReply)), + 'follow-up after save-and-submit branch', + ); + expectNoFoldedMessages(messages); + expectParent(messages, afterEditPrompt, editedMiddleReply, true); + expectParent(messages, afterEditReply, afterEditPrompt, false); + expect(messages.some((message) => messageText(message).includes(tailReply))).toBe(true); + expect(messages.some((message) => messageText(message).includes(regeneratedTailReply))).toBe( + true, + ); + + await reloadAndExpectMessages(page, [rootReply, firstReply, editedMiddleReply, afterEditReply]); + await revisitConversationAndExpectMessages(page, conversationId, [ + rootReply, + firstReply, + editedMiddleReply, + afterEditReply, + ]); + await clickSibling(page, editedMiddlePrompt, 'Previous'); + await expectVisibleMessages(page, [middlePrompt, fourthReply]); + await expectCanCycleSiblingTexts(page, tailReply, regeneratedTailReply); + await clickSibling(page, middlePrompt, 'Next'); + await expectVisibleMessages(page, [editedMiddlePrompt, editedMiddleReply, afterEditReply]); + }); + + test('error responses remain valid parents for follow-ups', async ({ page }) => { + const label = uniqueLabel('error'); + const basePrompt = replyPrompt(`${label}-base`); + const baseReply = replyText(`${label}-base`); + const errorPrompt = `E2E_FORCED_ERROR:${label}`; + const errorText = `E2E forced stream error ${label}`; + const afterErrorPrompt = replyPrompt(`${label}-after-error`); + const afterErrorReply = replyText(`${label}-after-error`); + + await openMockChat(page); + await sendAndExpectReply(page, basePrompt, baseReply); + const conversationId = await conversationIdFromPage(page); + + await sendAndExpectReply(page, errorPrompt, errorText); + await expect(messagesView(page).getByText(errorText)).toBeVisible({ timeout: 30000 }); + + await sendAndExpectReply(page, afterErrorPrompt, afterErrorReply); + const messages = await waitForMessages( + page, + conversationId, + (items) => items.some((message) => messageText(message).includes(afterErrorReply)), + 'follow-up after error', + ); + expectNoFoldedMessages(messages); + expectParent(messages, errorPrompt, baseReply, true); + expectParent(messages, errorText, errorPrompt, false); + expectParent(messages, afterErrorPrompt, errorText, true); + expectParent(messages, afterErrorReply, afterErrorPrompt, false); + + await reloadAndExpectMessages(page, [baseReply, errorText, afterErrorReply]); + await revisitConversationAndExpectMessages(page, conversationId, [ + baseReply, + errorText, + afterErrorReply, + ]); + }); + + test('generation-start failures recover without folding the next follow-up', async ({ page }) => { + const label = uniqueLabel('start-error'); + const basePrompt = replyPrompt(`${label}-base`); + const baseReply = replyText(`${label}-base`); + const failedPrompt = replyPrompt(`${label}-failed-start`); + const failedText = `E2E generation start failure ${label}`; + const afterFailurePrompt = replyPrompt(`${label}-after-start-failure`); + const afterFailureReply = replyText(`${label}-after-start-failure`); + + await openMockChat(page); + await sendAndExpectReply(page, basePrompt, baseReply); + const conversationId = await conversationIdFromPage(page); + + const failGenerationStart = async (route: Route) => { + const request = route.request(); + const { pathname } = new URL(request.url()); + const isAgentsChat = + pathname === '/api/agents/chat' || pathname.startsWith('/api/agents/chat/'); + if ( + request.method() !== 'POST' || + !isAgentsChat || + pathname.endsWith('/abort') || + !request.postData()?.includes(failedPrompt) + ) { + await route.continue(); + return; + } + + await route.fulfill({ + status: 500, + contentType: 'application/json', + body: JSON.stringify({ message: failedText }), + }); + }; + await page.route('**/api/agents/chat**', failGenerationStart); + + const failure = await submitMessageExpectingGenerationFailure(page, failedPrompt, 500); + expect(failure.ok()).toBe(false); + await expect(messagesView(page).getByText(failedText)).toBeVisible({ timeout: 30000 }); + await expect(page.getByRole('textbox', { name: 'Message input' })).toBeEnabled({ + timeout: 30000, + }); + await page.unroute('**/api/agents/chat**', failGenerationStart); + + await sendAndExpectReply(page, afterFailurePrompt, afterFailureReply); + const messages = await waitForMessages( + page, + conversationId, + (items) => items.some((message) => messageText(message).includes(afterFailureReply)), + 'follow-up after generation-start failure', + ); + expectNoFoldedMessages(messages); + expectParent(messages, afterFailurePrompt, baseReply, true); + expectParent(messages, afterFailureReply, afterFailurePrompt, false); + expect(messages.some((message) => messageText(message).includes(failedPrompt))).toBe(false); + expect(messages.some((message) => messageText(message).includes(failedText))).toBe(false); + + await reloadAndExpectMessages(page, [baseReply, afterFailureReply]); + await expect(messagesView(page).getByText(failedText)).toBeHidden(); + await revisitConversationAndExpectMessages(page, conversationId, [ + baseReply, + afterFailureReply, + ]); + await expect(messagesView(page).getByText(failedText)).toBeHidden(); + }); +});