diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index cb56912c60..ca27a838fe 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -1080,6 +1080,48 @@ describe('ResumableAgentController resume metadata', () => { ); }); + it('records regeneration ownership for exact-ID resume reconstruction', async () => { + const conversationId = 'conversation-123'; + const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Regenerate the edited response.', + messageId: 'user-message', + parentMessageId: 'parent-message', + responseMessageId: 'edited-response', + isRegenerate: true, + conversationId, + endpointOption: { + endpoint: 'agents', + modelOptions: { model: 'gpt-4.1' }, + }, + }, + config: {}, + }; + const res = { + headersSent: true, + json: jest.fn(() => { + res.headersSent = true; + }), + status: jest.fn(() => res), + }; + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( + conversationId, + 'user-123', + conversationId, + expect.objectContaining({ + initialMetadata: expect.objectContaining({ + responseMessageId: 'edited-response', + isRegenerate: true, + }), + }), + ); + }); + it('falls back to the model spec preset endpoint when no icon URL is configured', async () => { const conversationId = 'conversation-123'; const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index c6b9892e8f..1dedc6617d 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -1050,6 +1050,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // Persist temporary-chat state so a HITL resume keeps the resumed response // non-persisted instead of trusting the resume request to re-send the flag. isTemporary: req.body?.isTemporary, + ...(isRegenerate && { isRegenerate: true }), ...(scheduleId ? { scheduleId, diff --git a/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts b/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts index 49f5b37f14..08fb694367 100644 --- a/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts +++ b/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts @@ -4,6 +4,7 @@ import { buildCreatedInitialResponse, getExistingConversationAbortMessages, isInitialNewConversationSubmission, + mergeErrorMessages, mergeRegenerateFinalMessages, startedAsNewConversation, } from '~/hooks/SSE/useEventHandlers'; @@ -230,3 +231,61 @@ describe('getExistingConversationAbortMessages', () => { ).toEqual(['user-1']); }); }); + +describe('mergeErrorMessages', () => { + const message = (messageId: string, isCreatedByUser = false) => + ({ + messageId, + conversationId: 'conversation-1', + isCreatedByUser, + text: messageId, + }) as TMessage; + + it('adds the request and error for a normal submission', () => { + const userMessage = message('user-1', true); + const errorMessage = message('assistant-error'); + + expect( + mergeErrorMessages({ + messages: [message('previous-response')], + userMessage, + errorMessage, + }).map(({ messageId }) => messageId), + ).toEqual(['previous-response', 'user-1', 'assistant-error']); + }); + + it('preserves regeneration history without duplicating its user', () => { + const userMessage = message('user-1', true); + const originalResponse = message('assistant-1'); + const laterUser = message('user-2', true); + const laterResponse = message('assistant-2'); + const errorMessage = message('assistant-1_'); + + expect( + mergeErrorMessages({ + messages: [userMessage], + regenerateMessages: [userMessage, originalResponse, laterUser, laterResponse], + userMessage, + errorMessage, + isRegenerate: true, + }).map(({ messageId }) => messageId), + ).toEqual(['user-1', 'assistant-1', 'user-2', 'assistant-2', 'assistant-1_']); + }); + + it('replaces an edited response error that intentionally reuses its id', () => { + const userMessage = message('user-1', true); + const originalResponse = message('assistant-1'); + const errorMessage = { ...originalResponse, text: 'Regeneration failed', error: true }; + + const merged = mergeErrorMessages({ + messages: [userMessage], + regenerateMessages: [userMessage, originalResponse], + userMessage, + errorMessage, + isRegenerate: true, + }); + + expect(merged.map(({ messageId }) => messageId)).toEqual(['user-1', 'assistant-1']); + expect(merged[1]).toEqual(errorMessage); + }); +}); diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts index 4f09207441..49fd41fe8e 100644 --- a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts +++ b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts @@ -2964,14 +2964,19 @@ describe('useResumableSSE', () => { unmount(); }); - /** - * Regenerate: the new run's response id is not in the loaded history yet, so the - * parent-based fallback lands on the answer being REPLACED. Preserving that row's - * content would make the regenerated run's deltas append to the stale answer, so an - * empty snapshot must still clear a row we only matched heuristically. - */ - it('does not preserve content on a row matched only by the parent fallback', async () => { - const submission = buildSubmission(); + it('does not reuse an older response that only shares the user parent', async () => { + const submission = buildSubmission({ + initialResponse: { + messageId: 'resp-1', + conversationId: CONV_ID, + text: '', + isCreatedByUser: false, + sender: 'Custom Assistant', + endpoint: 'azureOpenAI', + iconURL: 'https://example.com/assistant.png', + model: 'gpt-4.1', + }, + }); const chatHelpers = buildChatHelpers(); chatHelpers.getMessages.mockReturnValue([ { @@ -3009,12 +3014,27 @@ describe('useResumableSSE', () => { }); }); - const synced = chatHelpers.setMessages.mock.calls + const syncedMessages = chatHelpers.setMessages.mock.calls .map(([messages]) => messages as TMessage[]) .reverse() - .find((messages) => messages?.some((m) => m.messageId === 'resp-previous')) - ?.find((m) => m.messageId === 'resp-previous'); - expect(synced?.content).toEqual([]); + .find((messages) => messages?.some((m) => m.messageId === 'resp-regenerated')); + expect(syncedMessages?.find((m) => m.messageId === 'resp-previous')?.content).toEqual([ + { type: 'text', text: 'the answer being regenerated' }, + ]); + expect(syncedMessages?.map((message) => message.messageId)).toEqual([ + 'msg-1', + 'resp-previous', + 'resp-regenerated', + ]); + expect(syncedMessages?.find((m) => m.messageId === 'resp-regenerated')).toEqual( + expect.objectContaining({ + content: [], + sender: 'Custom Assistant', + endpoint: 'azureOpenAI', + iconURL: 'https://example.com/assistant.png', + model: 'gpt-4.1', + }), + ); unmount(); }); @@ -4128,3 +4148,250 @@ describe('useResumableSSE', () => { unmount(); }); }); + +describe('useResumableSSE - sync response identity', () => { + beforeEach(() => { + mockSSEInstances.length = 0; + mockSetIsSubmitting.mockClear(); + }); + + const emitSync = async ( + sse: MockSSEInstance, + aggregatedContent: TMessage['content'], + responseMessageId?: string, + sender?: string, + userMessage?: Partial, + ) => { + await act(async () => { + sse._emit('message', { + data: JSON.stringify({ + sync: true, + resumeState: { aggregatedContent, responseMessageId, sender, userMessage }, + }), + }); + }); + }; + + it('updates the submission-owned response when sync omits the response ID', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const activeResponse = { + messageId: 'server-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: activeResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, activeResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + const aggregatedContent: TMessage['content'] = [ + { type: ContentTypes.TEXT, text: { value: 'Recovered answer' } }, + ]; + await emitSync(getLastSSE(), aggregatedContent); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages).toHaveLength(2); + expect(updatedMessages.find((message) => message.messageId === 'server-response-id')).toEqual({ + ...activeResponse, + content: aggregatedContent, + }); + expect( + updatedMessages.find((message) => message.messageId === 'server-user-id_'), + ).toBeUndefined(); + unmount(); + }); + + it('adds resumed sender metadata to an exact persisted response', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const activeResponse = { + messageId: 'server-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: activeResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, activeResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + await emitSync(getLastSSE(), [], activeResponse.messageId, 'Restored Assistant'); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages[1]).toEqual({ + ...activeResponse, + sender: 'Restored Assistant', + }); + unmount(); + }); + + it('appends a missing submission-owned response after older regeneration siblings', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const olderResponse = { + messageId: 'older-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Earlier answer', + content: [{ type: 'text', text: { value: 'Earlier answer' } }], + isCreatedByUser: false, + } as TMessage; + const activeResponse = { + messageId: 'active-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: activeResponse }), + isRegenerate: true, + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, olderResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + const aggregatedContent: TMessage['content'] = [ + { type: ContentTypes.TEXT, text: { value: 'Regenerated answer' } }, + ]; + await emitSync(getLastSSE(), aggregatedContent); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages.find((message) => message.messageId === 'older-response-id')).toEqual( + olderResponse, + ); + expect(updatedMessages.map((message) => message.messageId)).toEqual([ + 'server-user-id', + 'older-response-id', + 'active-response-id', + ]); + expect(updatedMessages.find((message) => message.messageId === 'active-response-id')).toEqual({ + ...activeResponse, + content: aggregatedContent, + }); + unmount(); + }); + + it('replaces the current-run placeholder without erasing its loaded content', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const preliminaryResponse = { + messageId: 'server-user-id_', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [{ type: ContentTypes.TEXT, text: { value: 'Already streaming' } }], + sender: 'Assistant', + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: preliminaryResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, preliminaryResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + await emitSync(getLastSSE(), [], 'assigned-response-id'); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages.map((message) => message.messageId)).toEqual([ + 'server-user-id', + 'assigned-response-id', + ]); + expect(updatedMessages[1]).toEqual({ + ...preliminaryResponse, + messageId: 'assigned-response-id', + }); + unmount(); + }); + + it('replaces the current-run user when sync assigns both durable IDs', async () => { + const preliminaryUser = { + messageId: 'client-user-id', + parentMessageId: 'previous-response-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const preliminaryResponse = { + messageId: 'client-user-id_', + parentMessageId: preliminaryUser.messageId, + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage: preliminaryUser, initialResponse: preliminaryResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([preliminaryUser, preliminaryResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + const assignedUser = { + ...preliminaryUser, + messageId: 'assigned-user-id', + }; + await emitSync(getLastSSE(), [], 'assigned-response-id', undefined, assignedUser); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages.map((message) => message.messageId)).toEqual([ + 'assigned-user-id', + 'assigned-response-id', + ]); + expect(updatedMessages[1]?.parentMessageId).toBe('assigned-user-id'); + unmount(); + }); +}); diff --git a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx index 93fb20a020..1af5beda97 100644 --- a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx +++ b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx @@ -290,6 +290,81 @@ describe('useResumeOnLoad', () => { expect(attached?.resumeGenerationCreatedAt).toBe(4242); }); + it('restores an externally started regeneration after history refreshes', async () => { + const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); + const olderResponse = { + messageId: 'older-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Older response', + isCreatedByUser: false, + } as TMessage; + const newerResponse = { + messageId: 'newer-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Newer response', + isCreatedByUser: false, + } as TMessage; + const observedSubmissions: Array = []; + let messages = [rootUser]; + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + + const { rerender, queryClient } = renderUseResumeOnLoad({ + getMessages: () => messages, + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + await act(async () => { + await Promise.resolve(); + }); + + const invalidate = jest.spyOn(queryClient, 'invalidateQueries').mockResolvedValue(undefined); + messages = [rootUser, newerResponse, olderResponse]; + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { + active: true, + status: 'running', + createdAt: 4242, + streamId: CONVERSATION_ID, + resumeState: { + aggregatedContent: [{ type: ContentTypes.TEXT, text: 'regenerating' }], + responseMessageId: `${olderResponse.messageId}_`, + userMessage: { + messageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + }, + }, + }, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(invalidate).toHaveBeenCalledWith({ + queryKey: [QueryKeys.messages, CONVERSATION_ID], + }); + const attached = observedSubmissions[observedSubmissions.length - 1]; + expect(attached?.isRegenerate).toBe(true); + expect(attached?.initialResponse?.messageId).toBe(`${olderResponse.messageId}_`); + expect(attached?.messages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + newerResponse.messageId, + olderResponse.messageId, + ]); + expect(attached?.regenerateMessages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + newerResponse.messageId, + olderResponse.messageId, + ]); + }); + it('re-arms once per run rather than on every poll of the active list', async () => { mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); const { rerender } = renderUseResumeOnLoad({ messages: [] }); @@ -771,6 +846,54 @@ describe('useResumeOnLoad', () => { ]); }); + it('does not claim an older sibling when resume state omits the response ID', async () => { + const observedSubmissions: Array = []; + const userMessage = buildUserMessage(CONVERSATION_ID); + const olderSibling = { + messageId: 'older-sibling-response', + parentMessageId: userMessage.messageId, + conversationId: CONVERSATION_ID, + text: 'Older sibling', + isCreatedByUser: false, + } as TMessage; + + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { + active: true, + status: 'running', + streamId: CONVERSATION_ID, + resumeState: { + runSteps: [], + aggregatedContent: [{ type: 'text', text: 'Active branch streaming' }], + conversationId: CONVERSATION_ID, + userMessage: { + messageId: userMessage.messageId, + parentMessageId: userMessage.parentMessageId, + conversationId: CONVERSATION_ID, + text: userMessage.text, + }, + }, + }, + }); + + renderUseResumeOnLoad({ + messages: [userMessage, olderSibling], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + + await act(async () => { + await Promise.resolve(); + }); + + const submission = observedSubmissions[observedSubmissions.length - 1]; + expect(submission?.initialResponse?.messageId).toBe(`${userMessage.messageId}_`); + expect((submission?.messages ?? []).map((message) => message.messageId)).toEqual([ + olderSibling.messageId, + ]); + }); + it('restores the branch that owns a pending OAuth resume user message', async () => { const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); const branchOneResponse = { @@ -834,13 +957,16 @@ describe('useResumeOnLoad', () => { expect(observedSiblingIndexes[observedSiblingIndexes.length - 1]).toBe(1); }); - it('restores the assistant sibling selected by a pending regenerate response', async () => { + it('restores the regenerate branch without claiming its older response', async () => { const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); const olderResponse = { messageId: 'older-response', parentMessageId: rootUser.messageId, conversationId: CONVERSATION_ID, text: 'Older response', + sender: 'Agent One', + model: 'gpt-5', + iconURL: 'https://example.com/agent-one.png', isCreatedByUser: false, } as TMessage; const newerResponse = { @@ -890,9 +1016,88 @@ describe('useResumeOnLoad', () => { expect(observedSiblingIndexes[observedSiblingIndexes.length - 1]).toBe(0); const submission = observedSubmissions[observedSubmissions.length - 1]; - expect(submission?.initialResponse?.messageId).toBe(olderResponse.messageId); + expect(submission?.initialResponse?.messageId).toBe(`${olderResponse.messageId}_`); + expect(submission?.initialResponse).toEqual( + expect.objectContaining({ + sender: olderResponse.sender, + model: olderResponse.model, + iconURL: olderResponse.iconURL, + }), + ); + expect(submission?.isRegenerate).toBe(true); expect((submission?.messages ?? []).map((message) => message.messageId)).toEqual([ + rootUser.messageId, newerResponse.messageId, + olderResponse.messageId, + ]); + expect((submission?.regenerateMessages ?? []).map((message) => message.messageId)).toEqual([ + rootUser.messageId, + newerResponse.messageId, + olderResponse.messageId, + ]); + }); + + it('preserves an exact-ID edited regeneration branch for early-abort rollback', async () => { + const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); + const editedResponse = { + messageId: 'edited-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Original response before the edit', + isCreatedByUser: false, + } as TMessage; + const siblingResponse = { + messageId: 'sibling-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Unrelated sibling', + isCreatedByUser: false, + } as TMessage; + const observedSubmissions: Array = []; + + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { + active: true, + status: 'running', + streamId: CONVERSATION_ID, + resumeState: { + runSteps: [], + aggregatedContent: [], + responseMessageId: editedResponse.messageId, + isRegenerate: true, + conversationId: CONVERSATION_ID, + userMessage: { + messageId: rootUser.messageId, + parentMessageId: rootUser.parentMessageId, + conversationId: CONVERSATION_ID, + text: rootUser.text, + }, + }, + }, + }); + + renderUseResumeOnLoad({ + messages: [rootUser, siblingResponse, editedResponse], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + + await act(async () => { + await Promise.resolve(); + }); + + const submission = observedSubmissions[observedSubmissions.length - 1]; + expect(submission?.isRegenerate).toBe(true); + expect(submission?.initialResponse?.messageId).toBe(editedResponse.messageId); + expect(submission?.messages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + siblingResponse.messageId, + ]); + expect(submission?.regenerateMessages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + siblingResponse.messageId, + editedResponse.messageId, ]); }); diff --git a/client/src/hooks/SSE/useEventHandlers.ts b/client/src/hooks/SSE/useEventHandlers.ts index 4c5b6b0e5f..570eafbcb5 100644 --- a/client/src/hooks/SSE/useEventHandlers.ts +++ b/client/src/hooks/SSE/useEventHandlers.ts @@ -185,6 +185,35 @@ export const getExistingConversationAbortMessages = ({ return [...sourceMessages]; }; +export const mergeErrorMessages = ({ + messages, + regenerateMessages, + userMessage, + errorMessage, + isRegenerate = false, +}: Pick & { + errorMessage: TMessage; +}): TMessage[] => { + if (isRegenerate) { + const finalMessages: TMessage[] = []; + let replaced = false; + for (const message of regenerateMessages ?? messages) { + if (message.messageId === errorMessage.messageId) { + finalMessages.push(errorMessage); + replaced = true; + } else { + finalMessages.push(message); + } + } + if (!replaced) { + finalMessages.push(errorMessage); + } + return finalMessages; + } + + return [...messages, userMessage, errorMessage]; +}; + export type EventHandlerParams = { isAddedRequest?: boolean; runIndex?: number; @@ -945,14 +974,14 @@ export default function useEventHandlers({ const errorHandler = useCallback( ({ data, submission }: { data?: TResData; submission: EventSubmission }) => { - const { messages, userMessage, initialResponse } = submission; + const { userMessage, initialResponse } = submission; setCompleted((prev) => new Set(prev.add(initialResponse.messageId))); const conversationId = userMessage.conversationId ?? submission.conversation?.conversationId ?? ''; const setErrorMessages = (convoId: string, errorMessage: TMessage) => { - const finalMessages: TMessage[] = [...messages, userMessage, errorMessage]; + const finalMessages = mergeErrorMessages({ ...submission, errorMessage }); setMessages(finalMessages); queryClient.setQueryData([QueryKeys.messages, convoId], finalMessages); }; diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 67ed895718..613a86e48c 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -583,18 +583,109 @@ const buildResumeEventSubmission = ( } as EventSubmission; }; +type ResumeMessageIndexes = { + userIndex: number; + responseIndex: number; + preliminaryUserIndex: number; + preliminaryResponseIndex: number; +}; + +const getResumeMessageIndexes = ( + messages: TMessage[], + userMessageId: string, + responseMessageId: string, + preliminaryUserMessageId?: string, + preliminaryResponseMessageId?: string, +): ResumeMessageIndexes => { + let userIndex = -1; + let responseIndex = -1; + let preliminaryUserIndex = -1; + let preliminaryResponseIndex = -1; + const hasPreliminaryResponse = preliminaryResponseMessageId?.endsWith('_') === true; + const eligiblePreliminaryUserId = + hasPreliminaryResponse && preliminaryUserMessageId && preliminaryUserMessageId !== userMessageId + ? preliminaryUserMessageId + : undefined; + const eligiblePreliminaryResponseId = + hasPreliminaryResponse && preliminaryResponseMessageId !== responseMessageId + ? preliminaryResponseMessageId + : undefined; + + for (let index = 0; index < messages.length; index++) { + const messageId = messages[index]?.messageId; + if (userIndex < 0 && messageId === userMessageId) { + userIndex = index; + } + if (responseIndex < 0 && messageId === responseMessageId) { + responseIndex = index; + } + if ( + preliminaryUserIndex < 0 && + eligiblePreliminaryUserId && + messageId === eligiblePreliminaryUserId + ) { + preliminaryUserIndex = index; + } + if ( + preliminaryResponseIndex < 0 && + eligiblePreliminaryResponseId && + messageId === eligiblePreliminaryResponseId + ) { + preliminaryResponseIndex = index; + } + } + + return { userIndex, responseIndex, preliminaryUserIndex, preliminaryResponseIndex }; +}; + const mergeResumeMessages = ( messages: TMessage[], userMessage: TMessage, responseMessage: TMessage, + indexes: ResumeMessageIndexes, ): TMessage[] => { const nextMessages = [...messages]; - const userIndex = nextMessages.findIndex( - (message) => message.messageId === userMessage.messageId, - ); - const responseIndex = nextMessages.findIndex( - (message) => message.messageId === responseMessage.messageId, - ); + let { userIndex, responseIndex, preliminaryResponseIndex } = indexes; + const { preliminaryUserIndex } = indexes; + + if (preliminaryUserIndex >= 0) { + if (userIndex >= 0) { + nextMessages.splice(preliminaryUserIndex, 1); + if (userIndex > preliminaryUserIndex) { + userIndex -= 1; + } + if (responseIndex > preliminaryUserIndex) { + responseIndex -= 1; + } + if (preliminaryResponseIndex > preliminaryUserIndex) { + preliminaryResponseIndex -= 1; + } + } else { + nextMessages[preliminaryUserIndex] = { + ...nextMessages[preliminaryUserIndex], + ...userMessage, + }; + userIndex = preliminaryUserIndex; + } + } + + if (preliminaryResponseIndex >= 0) { + if (responseIndex >= 0) { + nextMessages.splice(preliminaryResponseIndex, 1); + if (userIndex > preliminaryResponseIndex) { + userIndex -= 1; + } + if (responseIndex > preliminaryResponseIndex) { + responseIndex -= 1; + } + } else { + nextMessages[preliminaryResponseIndex] = { + ...nextMessages[preliminaryResponseIndex], + ...responseMessage, + }; + responseIndex = preliminaryResponseIndex; + } + } if (userIndex >= 0) { nextMessages[userIndex] = { ...nextMessages[userIndex], ...userMessage }; @@ -609,9 +700,7 @@ const mergeResumeMessages = ( } if (userIndex >= 0) { - const insertAt = userIndex + 1; - nextMessages.splice(insertAt, 0, responseMessage); - return nextMessages; + return [...nextMessages, responseMessage]; } if (responseIndex >= 0) { @@ -1848,6 +1937,16 @@ export default function useResumableSSE( const runId = v4(); setActiveRunId(runId); + /** Keep the current run's preliminary id long enough to replace that optimistic + * row in place if this snapshot assigns its durable response id. */ + const preliminaryResponseMessageId = currentSubmission.initialResponse?.messageId; + const currentUserMessageId = currentSubmission.userMessage?.messageId; + const preliminaryUserMessageId = + currentUserMessageId && + (currentSubmission.initialResponse?.parentMessageId === currentUserMessageId || + preliminaryResponseMessageId === `${currentUserMessageId}_`) + ? currentUserMessageId + : undefined; const resumeSubmission = buildResumeEventSubmission( currentSubmission, userMessage, @@ -1893,28 +1992,23 @@ export default function useResumableSSE( if (data.resumeState?.aggregatedContent && userMessage?.messageId) { const messages = getMessages() ?? []; const userMsgId = userMessage.messageId; - const serverResponseId = data.resumeState.responseMessageId; const hasResumedContent = data.resumeState.aggregatedContent.length > 0; - - let responseIdx = -1; - /** Only an id match proves the row belongs to THIS generation; the parent-based - * fallback below can land on a prior sibling (e.g. the answer being regenerated). */ - let matchedByResponseId = false; - if (serverResponseId) { - responseIdx = messages.findIndex((m) => m.messageId === serverResponseId); - matchedByResponseId = responseIdx >= 0; - } - if (responseIdx < 0) { - responseIdx = messages.findIndex( - (m) => - !m.isCreatedByUser && - (m.messageId === `${userMsgId}_` || m.parentMessageId === userMsgId), - ); - } + const responseId = resumeSubmission.initialResponse.messageId; + const messageIndexes = getResumeMessageIndexes( + messages, + userMsgId, + responseId, + preliminaryUserMessageId, + preliminaryResponseMessageId, + ); + const responseIdx = + messageIndexes.responseIndex >= 0 + ? messageIndexes.responseIndex + : messageIndexes.preliminaryResponseIndex; logger.log('ResumableSSE', 'SYNC update', { userMsgId, - serverResponseId, + responseId, responseIdx, foundMessageId: responseIdx >= 0 ? messages[responseIdx]?.messageId : null, messagesCount: messages.length, @@ -1924,15 +2018,11 @@ export default function useResumableSSE( if (responseIdx >= 0) { const oldContent = messages[responseIdx]?.content; /** An EMPTY resume snapshot is not authoritative over content we already loaded - * for the SAME response: assigning it would erase that content and leave a bare - * cursor. Restricted to an id match — preserving a fallback-matched row would - * make a regenerated run append to the answer it is replacing — and to a row - * that actually HAS parts, so the array is never swapped for `undefined`. */ + * for the SAME generation-owned response: assigning it would erase that content + * and leave a bare cursor. Require an existing content array so it is never + * swapped for `undefined`. */ const preserveLoadedContent = - !hasResumedContent && - matchedByResponseId && - Array.isArray(oldContent) && - oldContent.length > 0; + !hasResumedContent && Array.isArray(oldContent) && oldContent.length > 0; /** * Replacing the response with `aggregatedContent` drops the * prefix an edited resubmission had retained: the snapshot is @@ -1947,21 +2037,32 @@ export default function useResumableSSE( editPrefixClearedRef.current = true; } const responseMessage = { + ...resumeSubmission.initialResponse, ...messages[responseIdx], + messageId: responseId, + parentMessageId: userMsgId, content: preserveLoadedContent ? oldContent : data.resumeState.aggregatedContent, + sender: messages[responseIdx]?.sender ?? resumeSubmission.initialResponse.sender, iconURL: preferDefinedString( messages[responseIdx]?.iconURL, - data.resumeState.iconURL, + resumeSubmission.initialResponse.iconURL ?? data.resumeState.iconURL, + ), + model: preferDefinedString( + messages[responseIdx]?.model, + resumeSubmission.initialResponse.model ?? data.resumeState.model, ), - model: preferDefinedString(messages[responseIdx]?.model, data.resumeState.model), } as TMessage; - const updated = mergeResumeMessages(messages, userMessage, responseMessage); + const updated = mergeResumeMessages( + messages, + userMessage, + responseMessage, + messageIndexes, + ); logger.log('ResumableSSE', 'SYNC updating message', { messageId: responseMessage.messageId, oldContentLength: Array.isArray(oldContent) ? oldContent.length : 0, newContentLength: data.resumeState.aggregatedContent?.length, preservedExistingContent: preserveLoadedContent, - matchedByResponseId, }); setMessages(updated); resetContentHandler(); @@ -1975,18 +2076,14 @@ export default function useResumableSSE( * only in the matched branch left this path adding an offset * to indices that were already absolute. */ editPrefixClearedRef.current = true; - const responseId = serverResponseId ?? `${userMsgId}_`; const newMessage = { + ...resumeSubmission.initialResponse, messageId: responseId, parentMessageId: userMsgId, - conversationId: currentSubmission.conversation?.conversationId ?? '', - text: '', content: data.resumeState.aggregatedContent, isCreatedByUser: false, - iconURL: data.resumeState.iconURL, - model: data.resumeState.model, } as TMessage; - setMessages(mergeResumeMessages(messages, userMessage, newMessage)); + setMessages(mergeResumeMessages(messages, userMessage, newMessage, messageIndexes)); resetContentHandler(); syncStepMessage(newMessage); } diff --git a/client/src/hooks/SSE/useResumeOnLoad.ts b/client/src/hooks/SSE/useResumeOnLoad.ts index 9ff950cc73..f0b0f7653f 100644 --- a/client/src/hooks/SSE/useResumeOnLoad.ts +++ b/client/src/hooks/SSE/useResumeOnLoad.ts @@ -128,15 +128,27 @@ function buildSubmissionFromResumeState( (m) => m.isCreatedByUser && m.messageId === userMessageData?.messageId, ); - // Try to find existing response message in the messages array (from database). - // Regeneration can expose the in-flight placeholder id with trailing underscores - // while the persisted sibling uses the unpadded id. Prefer both exact identities - // before falling back to the shared parent, where several branch siblings can match. + // A trailing underscore distinguishes an in-flight regeneration from the persisted + // response it replaces. Only the exact response id proves generation ownership. + const existingResponseMessage = messages.find( + (m) => !m.isCreatedByUser && m.messageId === responseMessageId, + ); + // The persisted row may seed display metadata, but never identity or deduplication. const unpaddedResponseMessageId = responseMessageId.replace(/_+$/, ''); - const existingResponseMessage = - messages.find((m) => !m.isCreatedByUser && m.messageId === responseMessageId) ?? - messages.find((m) => !m.isCreatedByUser && m.messageId === unpaddedResponseMessageId) ?? - messages.find((m) => !m.isCreatedByUser && m.parentMessageId === userMessageData?.messageId); + const persistedRegenerationResponse = + unpaddedResponseMessageId !== responseMessageId + ? messages.find((m) => !m.isCreatedByUser && m.messageId === unpaddedResponseMessageId) + : undefined; + const responseMetadataMessage = existingResponseMessage ?? persistedRegenerationResponse; + const isRegenerateResume = + resumeState.isRegenerate === true || persistedRegenerationResponse != null; + let regenerateMessages: TMessage[] | undefined; + if (isRegenerateResume) { + regenerateMessages = + unpaddedResponseMessageId === responseMessageId + ? [...messages] + : messages.filter((message) => message.messageId !== responseMessageId); + } // Create or use existing user message const userMessage: TMessage = @@ -169,9 +181,9 @@ function buildSubmissionFromResumeState( content: (resumeState.aggregatedContent as TMessage['content']) ?? [], isCreatedByUser: false, role: 'assistant', - sender: existingResponseMessage?.sender ?? resumeState.sender, - model: preferDefinedString(existingResponseMessage?.model, resumeState.model), - iconURL: preferDefinedString(existingResponseMessage?.iconURL, resumeState.iconURL), + sender: responseMetadataMessage?.sender ?? resumeState.sender, + model: preferDefinedString(responseMetadataMessage?.model, resumeState.model), + iconURL: preferDefinedString(responseMetadataMessage?.iconURL, resumeState.iconURL), } as TMessage; // Re-paused turn: seed the approval / ask-user controls straight onto the @@ -186,17 +198,13 @@ function buildSubmissionFromResumeState( endpoint: null, } as TConversation; - // On reload, `messages` is the full DB array, which already holds the paused user - // row and the partial (unfinished) assistant row under the same ids that - // `userMessage` / `initialResponse` (and the resume final event's request/response - // messages) re-supply. Strip them so createdHandler/finalHandler — which build - // `[...messages, requestMessage, responseMessage]` — don't append a duplicate pair. - const pausedResponseIdUnpadded = initialResponse.messageId.replace(/_+$/, ''); + // Non-regenerate resumes strip the persisted request/response pair before handlers + // re-supply it. A regeneration keeps the original branch for early-abort rollback; + // explicit resume metadata covers edited regenerations that reuse the exact response id. const dedupedMessages = messages.filter( (m) => - m.messageId !== userMessage.messageId && m.messageId !== initialResponse.messageId && - m.messageId !== pausedResponseIdUnpadded, + (isRegenerateResume || m.messageId !== userMessage.messageId), ); return { @@ -204,7 +212,8 @@ function buildSubmissionFromResumeState( userMessage, initialResponse, conversation, - isRegenerate: false, + isRegenerate: isRegenerateResume, + ...(regenerateMessages && { regenerateMessages }), isTemporary: false, endpointOption: {}, // Signal to useResumableSSE to subscribe to existing stream instead of starting new diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 3a89dfb2d9..25febf31ab 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -2528,6 +2528,7 @@ class GenerationJobManagerClass { generationProtocolVersion: jobData.generationProtocolVersion, userMessage: jobData.userMessage, responseMessageId: jobData.responseMessageId, + isRegenerate: jobData.isRegenerate, sender: jobData.sender, endpoint: jobData.endpoint, iconURL: jobData.iconURL, @@ -6928,6 +6929,7 @@ class GenerationJobManagerClass { aggregatedContent, userMessage: jobData.userMessage, responseMessageId: jobData.responseMessageId, + isRegenerate: jobData.isRegenerate, conversationId: jobData.conversationId, sender: jobData.sender, iconURL: jobData.iconURL, diff --git a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts index d252af0d82..dee6a9e117 100644 --- a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts +++ b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts @@ -67,6 +67,22 @@ describe('GenerationJobManager resume replay events', () => { manager = undefined; }); + test('projects regeneration ownership into resume state', async () => { + manager = createInMemoryManager(); + const streamId = `regenerate-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId, { + initialMetadata: { + responseMessageId: 'edited-response', + isRegenerate: true, + }, + }); + + await expect(manager.getResumeState(streamId)).resolves.toMatchObject({ + responseMessageId: 'edited-response', + isRegenerate: true, + }); + }); + test('includes OAuth run step and delta replay events in resume state', async () => { manager = createInMemoryManager(); const streamId = `oauth-delta-resume-${Date.now()}`; diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts index 7298dd4026..e114c8d3c5 100644 --- a/packages/api/src/stream/__tests__/startup.spec.ts +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -99,6 +99,7 @@ describe('GenerationJobManager startup telemetry', () => { parentMessageId: 'parent-1', }, responseMessageId: 'response-1', + isRegenerate: true, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', @@ -129,6 +130,7 @@ describe('GenerationJobManager startup telemetry', () => { parentMessageId: 'parent-1', }, responseMessageId: 'response-1', + isRegenerate: true, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index ea46b4a4e1..b73a4d7af9 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -4547,6 +4547,7 @@ export class RedisJobStore implements IJobStoreV2 { recoveredSteerId: data.recoveredSteerId || undefined, userMessage: data.userMessage ? JSON.parse(data.userMessage) : undefined, responseMessageId: data.responseMessageId || undefined, + isRegenerate: data.isRegenerate != null ? data.isRegenerate === '1' : undefined, createdEventEmitted: data.createdEventEmitted === '1', sender: data.sender || undefined, syncSent: data.syncSent === '1', diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index a03e28298d..e086b51760 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -75,6 +75,9 @@ export interface SerializableJobData { /** Response message ID for reconnection */ responseMessageId?: string; + /** Whether this generation replaces an existing assistant branch. */ + isRegenerate?: boolean; + /** * Whether this run has activity labels enabled (per-endpoint * `activityLabel: true`). Set once at run start so the resume path can @@ -294,6 +297,7 @@ export type JobMetadataPatch = Partial< Pick< SerializableJobData, | 'responseMessageId' + | 'isRegenerate' | 'sender' | 'conversationId' | 'userMessage' diff --git a/packages/api/src/stream/metadata.ts b/packages/api/src/stream/metadata.ts index b334412267..da8c108b34 100644 --- a/packages/api/src/stream/metadata.ts +++ b/packages/api/src/stream/metadata.ts @@ -6,6 +6,9 @@ export function sanitizeJobMetadata(metadata: Partial): J if (metadata.responseMessageId) { patch.responseMessageId = metadata.responseMessageId; } + if (metadata.isRegenerate !== undefined) { + patch.isRegenerate = metadata.isRegenerate; + } if (metadata.sender) { patch.sender = metadata.sender; } diff --git a/packages/api/src/types/stream.ts b/packages/api/src/types/stream.ts index b6727ebfbf..fca5d8e943 100644 --- a/packages/api/src/types/stream.ts +++ b/packages/api/src/types/stream.ts @@ -17,6 +17,8 @@ export interface GenerationJobMetadata { userMessage?: Agents.UserMessageMeta; /** Response message ID for tracking */ responseMessageId?: string; + /** Whether this generation replaces an existing assistant branch. */ + isRegenerate?: boolean; /** Sender label for the response (e.g., "GPT-4.1", "Claude") */ sender?: string; /** Endpoint identifier for abort handling */ diff --git a/packages/data-provider/src/types/agents.ts b/packages/data-provider/src/types/agents.ts index ae8499c61f..35d2e39c9f 100644 --- a/packages/data-provider/src/types/agents.ts +++ b/packages/data-provider/src/types/agents.ts @@ -263,6 +263,8 @@ export namespace Agents { aggregatedContent?: MessageContentComplex[]; userMessage?: UserMessageMeta; responseMessageId?: string; + /** True when the live generation replaces an existing assistant branch. */ + isRegenerate?: boolean; conversationId?: string; sender?: string; iconURL?: string;