🎯 fix: Keep Run Steps and Labels in One Index Space After a Resume Sync (#14516)

Follow-up to #14391, which deliberately left this shared math untouched.

An edited resubmission offsets incoming indices past the prefix the client
retained, because the server indexes only NEW content. A resume sync
invalidates that arrangement twice, and run steps honored neither:

- It REPLACES `initialResponse.content` with the server's completion-local
  snapshot, so the live array stops measuring the retained prefix. Run
  steps derived their offset from that array, so a reconnect that produced
  an empty snapshot silently dropped the offset to zero and wrote over
  retained content.
- When it also replaces the RENDERED content, the prefix is gone entirely
  and server indices are already absolute. Run steps kept adding the
  snapshot's own length on top, writing past the end and leaving holes.

Activity labels already honored both facts (`editPrefixLength` +
`editPrefixClearedRef`), so a batch's tool cards and its header could
resolve in different index spaces: a label overwriting an unrelated part,
or a fill missing its own reservation and leaving the placeholder pending
forever.

Run steps now read the same two inputs. `useStepHandler` takes the
CAPTURED `editPrefixLength` rather than measuring the live array, gated on
a new `editPrefixCleared` flag that the resumable transport — which owns
the sync boundary — stamps onto dispatched submissions. The non-resumable
transport never sets it, so the plain edit path is unchanged.
`calculateContentIndex` now takes the offset directly instead of the
prefix array, so its ±1 trailing-text adjustment cannot diverge from the
offset every other path applies.

Tests (useStepHandler.spec): unedited applies no offset; a plain edit
still offsets; the captured length wins when sync replaced the live array;
a cleared prefix stops offsetting for both run steps and message deltas,
staying at absolute indices. Verified against the pre-fix code — the three
states this PR repairs fail there and pass here.
This commit is contained in:
Danny Avila 2026-07-29 15:39:10 -04:00 committed by GitHub
parent becfc5a373
commit d70cab48fd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 282 additions and 21 deletions

View file

@ -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);
});
});
});

View file

@ -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<typeof rawStepHandler>) =>
rawStepHandler(
event,
editPrefixClearedRef.current
? ({ ...submission, editPrefixCleared: true } as EventSubmission)
: submission,
),
[rawStepHandler],
);
const { data: startupConfig } = useGetStartupConfig();
const balanceQuery = useGetUserBalance({

View file

@ -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,

View file

@ -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. */