diff --git a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts index ac8c7c6240..f8455d5068 100644 --- a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts +++ b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts @@ -3135,4 +3135,208 @@ describe('useStepHandler', () => { expect(getProgress('call_keep')).not.toBeNull(); }); }); + + /** + * Edited resubmissions offset every incoming index past the prefix the + * client retained, because the server indexes only the NEW content. A + * resume SYNC invalidates that arrangement twice over: it REPLACES + * `initialResponse.content` with the server's completion-local snapshot + * (so the live array no longer measures the prefix), and when it also + * replaces the RENDERED content the prefix is gone entirely and server + * indices become absolute. + * + * Activity labels already honored both facts; run steps did not, so a + * batch's tool cards and its header could resolve in different index + * spaces. These cases pin the states no other suite constructs. + */ + describe('edit-prefix index space across resume', () => { + const textPart = (text: string): TMessageContentParts => + ({ type: ContentTypes.TEXT, [ContentTypes.TEXT]: text }) as TMessageContentParts; + + /** A retained non-text part, so the trailing-text `-1` adjustment stays out + * of offset assertions. */ + const keptToolPart = (): TMessageContentParts => + ({ + type: ContentTypes.TOOL_CALL, + [ContentTypes.TOOL_CALL]: { + id: 'kept-tool', + name: 'kept_tool', + args: '{}', + type: ToolCallTypes.TOOL_CALL, + }, + }) as unknown as TMessageContentParts; + + /** Renders the hook against a live message array and returns the response. */ + const runToolStepAt = (serverIndex: number, submission: EventSubmission) => { + const responseMessage = submission.initialResponse as TMessage; + let currentMessages: TMessage[] = [responseMessage]; + mockGetMessages.mockImplementation(() => currentMessages); + mockSetMessages.mockImplementation((messages: TMessage[]) => { + currentMessages = messages; + }); + + const { result } = renderHook(() => useStepHandler(createHookParams())); + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP, + data: createToolCallRunStep({ index: serverIndex, runId: responseMessage.messageId }), + }, + submission, + ); + }); + return currentMessages.find((m) => !m.isCreatedByUser); + }; + + it('applies no offset for an unedited submission', () => { + const response = runToolStepAt(0, createSubmission()); + expect(getToolCallName(response?.content?.[0])).toBe('test_tool'); + }); + + it('offsets by the retained prefix for an edited submission', () => { + const submission = createSubmission({ + editedContent: { index: 0, type: ContentTypes.TEXT }, + initialResponse: createResponseMessage({ + content: [textPart('kept a'), textPart('kept b')], + }), + } as never); + (submission as { editPrefixLength?: number }).editPrefixLength = 2; + + const response = runToolStepAt(0, submission); + /** Server index 0 is the first NEW part, so it lands past the prefix. */ + expect(getToolCallName(response?.content?.[2])).toBe('test_tool'); + expect(response?.content?.[0]).toMatchObject({ [ContentTypes.TEXT]: 'kept a' }); + }); + + /** + * The `preserveLoadedContent` shape: SYNC replaced `initialResponse.content` + * with an empty completion-local snapshot but KEPT the rendered prefix, so + * the flag stays clear. Measuring the live array would drop the offset to + * zero and overwrite retained content. + */ + it('offsets by the captured length when sync replaced the live content array', () => { + const submission = createSubmission({ + editedContent: { index: 0, type: ContentTypes.TEXT }, + initialResponse: createResponseMessage({ content: [] }), + } as never); + (submission as { editPrefixLength?: number }).editPrefixLength = 2; + + const response = runToolStepAt(0, submission); + expect(getToolCallName(response?.content?.[2])).toBe('test_tool'); + }); + + /** + * Post-SYNC: the rendered content IS the completion-local snapshot, so + * incoming indices are already absolute — the same space activity labels + * switch to when the flag is set. Any offset here writes past the end. + */ + it('applies no offset once the prefix is cleared by a resume sync', () => { + const submission = createSubmission({ + editedContent: { index: 0, type: ContentTypes.TEXT }, + initialResponse: createResponseMessage({ content: [] }), + } as never); + Object.assign(submission, { editPrefixLength: 2, editPrefixCleared: true }); + + const response = runToolStepAt(0, submission); + /** Offset still applied, this would land at index 2 and leave a hole. */ + expect(getToolCallName(response?.content?.[0])).toBe('test_tool'); + expect(response?.content).toHaveLength(1); + }); + + it('keeps a cleared-prefix tool card at its absolute server index', () => { + const submission = createSubmission({ + editedContent: { index: 0, type: ContentTypes.TEXT }, + initialResponse: createResponseMessage({ + content: [textPart('snapshot a'), textPart('snapshot b')], + }), + } as never); + Object.assign(submission, { editPrefixLength: 2, editPrefixCleared: true }); + + const response = runToolStepAt(2, submission); + /** Absolute: index 2 stays 2, appending after the snapshot's own parts. */ + expect(getToolCallName(response?.content?.[2])).toBe('test_tool'); + expect(response?.content?.[0]).toMatchObject({ [ContentTypes.TEXT]: 'snapshot a' }); + }); + + /** + * Deltas resolve through `calculateContentIndex`, which now receives the + * offset directly instead of measuring the prefix array. A prefix ending + * in a NON-text part isolates the offset from the separate `-1` + * adjustment that intentionally continues a retained trailing text part. + */ + it('offsets message deltas by the same prefix as run steps', () => { + const submission = createSubmission({ + editedContent: { index: 0, type: ContentTypes.TEXT }, + initialResponse: createResponseMessage({ + content: [textPart('kept a'), keptToolPart()], + }), + } as never); + (submission as { editPrefixLength?: number }).editPrefixLength = 2; + + const responseMessage = submission.initialResponse as TMessage; + let currentMessages: TMessage[] = [responseMessage]; + mockGetMessages.mockImplementation(() => currentMessages); + mockSetMessages.mockImplementation((messages: TMessage[]) => { + currentMessages = messages; + }); + + const { result } = renderHook(() => useStepHandler(createHookParams())); + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP, + data: createRunStep({ index: 0, runId: responseMessage.messageId }), + }, + submission, + ); + result.current.stepHandler( + { event: StepEvents.ON_MESSAGE_DELTA, data: createMessageDelta('step-1', 'streamed') }, + submission, + ); + }); + + const response = currentMessages.find((m) => !m.isCreatedByUser); + expect(response?.content?.[2]).toMatchObject({ [ContentTypes.TEXT]: 'streamed' }); + expect(response?.content?.[0]).toMatchObject({ [ContentTypes.TEXT]: 'kept a' }); + }); + + /** + * Post-sync a delta continues at the snapshot's NEXT absolute slot. With + * the offset still applied it would jump past that slot and leave a hole + * the label reservation could never line up with. + */ + it('applies no delta offset once the prefix is cleared', () => { + const submission = createSubmission({ + editedContent: { index: 0, type: ContentTypes.TEXT }, + initialResponse: createResponseMessage({ content: [keptToolPart(), keptToolPart()] }), + } as never); + Object.assign(submission, { editPrefixLength: 2, editPrefixCleared: true }); + + const responseMessage = submission.initialResponse as TMessage; + let currentMessages: TMessage[] = [responseMessage]; + mockGetMessages.mockImplementation(() => currentMessages); + mockSetMessages.mockImplementation((messages: TMessage[]) => { + currentMessages = messages; + }); + + const { result } = renderHook(() => useStepHandler(createHookParams())); + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP, + data: createRunStep({ index: 2, runId: responseMessage.messageId }), + }, + submission, + ); + result.current.stepHandler( + { event: StepEvents.ON_MESSAGE_DELTA, data: createMessageDelta('step-1', 'streamed') }, + submission, + ); + }); + + const response = currentMessages.find((m) => !m.isCreatedByUser); + expect(response?.content?.[2]).toMatchObject({ [ContentTypes.TEXT]: 'streamed' }); + expect(response?.content).toHaveLength(3); + }); + }); }); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index b0cebbd71d..0701c35135 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -666,10 +666,28 @@ export default function useResumableSSE( setShowStopButton, }); - /** Run steps dispatch straight through: their index math is upstream's and - * is deliberately left untouched by this feature. Only the activity-label - * handler applies the resume-aware prefix offset. */ - const stepHandler = rawStepHandler; + /** + * Run steps and activity labels must resolve indices in ONE space, and the + * resume SYNC boundary that invalidates the edit prefix is owned here — so + * this transport stamps its cleared state onto every dispatched + * submission. `useStepHandler` reads the flag (alongside the captured + * `editPrefixLength`) instead of measuring the live + * `initialResponse.content`, which SYNC replaces with the server's + * completion-local snapshot. + * + * Only allocates once the prefix is actually cleared; before that, and on + * the non-resumable transport, the submission passes through untouched. + */ + const stepHandler = useCallback( + (...[event, submission]: Parameters) => + rawStepHandler( + event, + editPrefixClearedRef.current + ? ({ ...submission, editPrefixCleared: true } as EventSubmission) + : submission, + ), + [rawStepHandler], + ); const { data: startupConfig } = useGetStartupConfig(); const balanceQuery = useGetUserBalance({ diff --git a/client/src/hooks/SSE/useStepHandler.ts b/client/src/hooks/SSE/useStepHandler.ts index 0055ee9664..5b41e88373 100644 --- a/client/src/hooks/SSE/useStepHandler.ts +++ b/client/src/hooks/SSE/useStepHandler.ts @@ -331,27 +331,31 @@ export default function useStepHandler({ /** * Calculate content index for a run step. - * For edited content scenarios, offset by initialContent length. + * + * Takes the edit-prefix OFFSET rather than the prefix array: after a resume + * sync the live array no longer describes the retained prefix, so deriving + * the offset here from its length would disagree with the offset every + * other event path applies. */ const calculateContentIndex = useCallback( ( serverIndex: number, - initialContent: TMessageContentParts[], + editPrefixOffset: number, incomingContentType: string, existingContent?: TMessageContentParts[], ): number => { /** Only apply -1 adjustment for TEXT or THINK types when they match existing content */ if ( - initialContent.length > 0 && + editPrefixOffset > 0 && (incomingContentType === ContentTypes.TEXT || incomingContentType === ContentTypes.THINK) ) { - const targetIndex = serverIndex + initialContent.length - 1; + const targetIndex = serverIndex + editPrefixOffset - 1; const existingType = existingContent?.[targetIndex]?.type; if (existingType === incomingContentType) { return targetIndex; } } - return serverIndex + initialContent.length; + return serverIndex + editPrefixOffset; }, [], ); @@ -696,10 +700,33 @@ export default function useStepHandler({ lastAnnouncementTimeRef.current = currentTime; } + /** + * Index offset for an edited resubmission: the server indexes only the + * NEW content, so incoming indices shift past the prefix the client + * kept. + * + * Reads the length CAPTURED when the submission was built rather than + * the live `initialResponse.content` array, because a resume sync + * REPLACES that array with the server's completion-local snapshot — + * whose length describes the new generation, not the retained prefix. + * They are equal until a reconnect, so the non-resumed path is + * unaffected. + * + * `editPrefixCleared` means that sync also replaced the RENDERED + * content: the prefix is gone from the message and server indices are + * already absolute, so any offset would write past the end. Activity + * labels honor the same flag — both must agree, or a batch's tool + * cards and its header land in different index spaces. + * + * `initialContent` stays the live array: it seeds a response that is + * not in the map yet, and post-sync the seeding path correctly falls + * back to the rendered content instead. + */ let initialContent: TMessageContentParts[] = []; - // For editedContent scenarios, use the initial response content for index offsetting - if (submission?.editedContent != null) { + let editPrefixOffset = 0; + if (submission?.editedContent != null && submission?.editPrefixCleared !== true) { initialContent = submission?.initialResponse?.content ?? initialContent; + editPrefixOffset = submission?.editPrefixLength ?? initialContent.length; } if (stepEvent.event === StepEvents.ON_RUN_STEP) { @@ -716,8 +743,8 @@ export default function useStepHandler({ stepMap.current.set(runStep.id, runStep); - // Calculate content index - use server index, offset by initialContent for edit scenarios - const contentIndex = runStep.index + initialContent.length; + // Calculate content index - use server index, offset by the retained edit prefix + const contentIndex = runStep.index + editPrefixOffset; let response = messageMap.current.get(responseMessageId); @@ -857,7 +884,7 @@ export default function useStepHandler({ const response = messageMap.current.get(responseMessageId); if (response) { // Agent updates don't need index adjustment - const currentIndex = agent_update.index + initialContent.length; + const currentIndex = agent_update.index + editPrefixOffset; // Agent updates carry their own agentId - use default groupId if agentId is present const agentUpdateMeta: ContentMetadata | undefined = agent_update.agentId ? { agentId: agent_update.agentId, groupId: 1 } @@ -909,7 +936,7 @@ export default function useStepHandler({ } const currentIndex = calculateContentIndex( runStep.index, - initialContent, + editPrefixOffset, contentPart.type || '', updatedResponse.content, ); @@ -959,7 +986,7 @@ export default function useStepHandler({ } const currentIndex = calculateContentIndex( runStep.index, - initialContent, + editPrefixOffset, contentPart.type || '', updatedResponse.content, ); @@ -1018,8 +1045,8 @@ export default function useStepHandler({ contentPart.tool_call.expires_at = runStepDelta.delta.expires_at; } - // Use server's index, offset by initialContent for edit scenarios - const currentIndex = runStep.index + initialContent.length; + // Use server's index, offset by the retained edit prefix + const currentIndex = runStep.index + editPrefixOffset; updatedResponse = updateContent( updatedResponse, currentIndex, @@ -1067,8 +1094,8 @@ export default function useStepHandler({ tool_call: result.tool_call, }; - // Use server's index, offset by initialContent for edit scenarios - const currentIndex = runStep.index + initialContent.length; + // Use server's index, offset by the retained edit prefix + const currentIndex = runStep.index + editPrefixOffset; updatedResponse = updateContent( updatedResponse, currentIndex, @@ -1113,7 +1140,7 @@ export default function useStepHandler({ summarizing: true, }; - const contentIndex = runStep.index + initialContent.length; + const contentIndex = runStep.index + editPrefixOffset; const updatedResponse = updateContent( response, contentIndex, diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index 7b00aeee59..524ce83d07 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -187,6 +187,18 @@ export type TSubmission = { * resumes for run steps and activity labels alike. */ editPrefixLength?: number; + /** + * Set once a resume SYNC has replaced the response's retained prefix with + * the server's completion-local snapshot. From that point the prefix is + * gone from the rendered message and server indices are absolute in the + * new space, so {@link editPrefixLength} must NOT be applied — by run + * steps or by activity labels. Both event paths read this flag so a + * batch's tool cards and its header always land in one index space. + * + * Stamped per event by the resumable transport, which owns the SYNC + * boundary; the non-resumable path never sets it. + */ + editPrefixCleared?: boolean; /** Added conversation for multi-convo feature */ addedConvo?: TConversation; /** Skills the user invoked via the `$` popover for this submission. */