diff --git a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts index 310f9bd0d8..08a6723b67 100644 --- a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts +++ b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts @@ -2704,7 +2704,7 @@ describe('useStepHandler', () => { expect(summaryPart).toMatchObject({ summarizing: false }); }); - it('ON_SUMMARIZE_COMPLETE error removes summarizing parts', () => { + it('ON_SUMMARIZE_COMPLETE error finalizes the part in place without splicing it out', () => { mockLastAnnouncementTimeRef.current = Date.now(); const responseMessage = createResponseMessage(); mockGetMessages.mockReturnValue([responseMessage]); @@ -2775,11 +2775,21 @@ describe('useStepHandler', () => { expect(mockSetMessages).toHaveBeenCalled(); const lastCall = mockSetMessages.mock.calls[mockSetMessages.mock.calls.length - 1][0]; const responseMsg = lastCall.find((m: TMessage) => m.messageId === 'response-msg-1'); + /** + * Failed rounds must keep their slot: splicing shifts every later part + * under the index-keyed renderer and breaks the position == step-index + * invariant that updateContent writes rely on. The part is finalized in + * place with its streamed content preserved. + */ const summaryParts = responseMsg?.content?.filter( (c: TMessageContentParts) => c.type === ContentTypes.SUMMARY, ) ?? []; - expect(summaryParts).toHaveLength(0); + expect(summaryParts).toHaveLength(1); + expect((summaryParts[0] as SummaryContentPart).summarizing).toBe(false); + expect((summaryParts[0] as SummaryContentPart).content).toEqual([ + { type: ContentTypes.TEXT, text: 'partial' }, + ]); }); it('ON_SUMMARIZE_COMPLETE returns early when target message not in messageMap', () => { diff --git a/client/src/hooks/SSE/useStepHandler.ts b/client/src/hooks/SSE/useStepHandler.ts index 934fad74a0..8abd553b16 100644 --- a/client/src/hooks/SSE/useStepHandler.ts +++ b/client/src/hooks/SSE/useStepHandler.ts @@ -1437,37 +1437,47 @@ export default function useStepHandler({ return; } - if (completeData.error) { - const filtered = targetMessage.content.filter( - (part) => - part?.type !== ContentTypes.SUMMARY || !(part as SummaryContentPart).summarizing, - ); - if (filtered.length !== targetMessage.content.length) { - announcePolite({ message: 'summarize_failed', isStatus: true }); - const cleaned = { ...targetMessage, content: filtered }; - const currentMessages = submission.isRegenerate ? messages : getMessages() || []; - messageMap.current.set(completeMessageId, cleaned); - setMessages(mergeResponseMessage(currentMessages, cleaned, completeMessageId)); - } - } else { - let didFinalize = false; - const updatedContent = targetMessage.content.map((part) => { - if (part?.type === ContentTypes.SUMMARY && (part as SummaryContentPart).summarizing) { - didFinalize = true; - if (!completeData.summary) { - return { ...part, summarizing: false } as SummaryContentPart; - } - return { ...completeData.summary, summarizing: false } as SummaryContentPart; - } + /** + * Scoped to the owning step's slot when the step is known: a global + * scan finalizes a NEWER round's in-flight part when summarize + * cycles run back-to-back (tiny context windows re-trigger + * summarization every graph step). Unknown step falls back to + * finalizing every in-flight part. + */ + const completeIndex = + completeRunStep != null ? completeRunStep.index + editPrefixOffset : -1; + let didFinalize = false; + const updatedContent = targetMessage.content.map((part, index) => { + if (part?.type !== ContentTypes.SUMMARY || !(part as SummaryContentPart).summarizing) { return part; - }); - if (didFinalize) { - announcePolite({ message: 'summarize_completed', isStatus: true }); - const finalized = { ...targetMessage, content: updatedContent }; - const currentMessages = submission.isRegenerate ? messages : getMessages() || []; - messageMap.current.set(completeMessageId, finalized); - setMessages(mergeResponseMessage(currentMessages, finalized, completeMessageId)); } + if (completeIndex >= 0 && index !== completeIndex) { + return part; + } + didFinalize = true; + if (!completeData.error && completeData.summary) { + return { ...completeData.summary, summarizing: false } as SummaryContentPart; + } + /** + * Failed rounds keep their slot. Splicing the part out shifts every + * later part under the index-keyed renderer (remounting rows and + * collapsing whatever the user expanded mid-stream) and breaks the + * content-position == step-index invariant that `updateContent` + * writes rely on. Flipping the flag alone hides an empty row + * (`Summary` renders null without text) and matches the persisted + * message, which retains the part server-side. + */ + return { ...part, summarizing: false } as SummaryContentPart; + }); + if (didFinalize) { + announcePolite({ + message: completeData.error ? 'summarize_failed' : 'summarize_completed', + isStatus: true, + }); + const finalized = { ...targetMessage, content: updatedContent }; + const currentMessages = submission.isRegenerate ? messages : getMessages() || []; + messageMap.current.set(completeMessageId, finalized); + setMessages(mergeResponseMessage(currentMessages, finalized, completeMessageId)); } } else { const _exhaustive: never = stepEvent;