From f77eac15bfe4b557026a83e8cc18fd4c41c9100b Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Sun, 9 Aug 2026 20:47:06 -0700 Subject: [PATCH 1/2] fix(workflow): invalidate UI parts when a step retries --- .changeset/fix-workflow-retry-invalidation.md | 6 ++ .../src/ui/process-ui-message-stream.test.ts | 62 +++++++++++++ .../ai/src/ui/process-ui-message-stream.ts | 26 +++++- packages/workflow/src/do-stream-step.test.ts | 79 ++++++++++++++++ packages/workflow/src/do-stream-step.ts | 89 +++++++++++++------ .../workflow/src/to-ui-message-chunk.test.ts | 17 ++++ packages/workflow/src/to-ui-message-chunk.ts | 7 ++ 7 files changed, 256 insertions(+), 30 deletions(-) create mode 100644 .changeset/fix-workflow-retry-invalidation.md create mode 100644 packages/workflow/src/do-stream-step.test.ts create mode 100644 packages/workflow/src/to-ui-message-chunk.test.ts diff --git a/.changeset/fix-workflow-retry-invalidation.md b/.changeset/fix-workflow-retry-invalidation.md new file mode 100644 index 000000000000..24b9820adaf8 --- /dev/null +++ b/.changeset/fix-workflow-retry-invalidation.md @@ -0,0 +1,6 @@ +--- +'ai': patch +'@ai-sdk/workflow': patch +--- + +Invalidate streamed UI message parts when a workflow model step is retried. diff --git a/packages/ai/src/ui/process-ui-message-stream.test.ts b/packages/ai/src/ui/process-ui-message-stream.test.ts index b9e537d30310..d964d0ddde24 100644 --- a/packages/ai/src/ui/process-ui-message-stream.test.ts +++ b/packages/ai/src/ui/process-ui-message-stream.test.ts @@ -4769,6 +4769,68 @@ describe('processUIMessageStream', () => { }); }); + it('should discard the current step when receiving a reload marker', async () => { + const stream = createUIMessageStream([ + { type: 'start', messageId: 'msg-123' }, + { type: 'start-step' }, + { + type: 'tool-input-available', + toolCallId: 'old-tool-call', + toolName: 'deleteFile', + input: { path: '/tmp/old.txt' }, + }, + { + type: 'tool-approval-request', + approvalId: 'old-approval', + toolCallId: 'old-tool-call', + }, + { type: 'finish-step' }, + { type: 'data-reload', data: {}, transient: true }, + { type: 'start-step' }, + { + type: 'tool-input-available', + toolCallId: 'new-tool-call', + toolName: 'deleteFile', + input: { path: '/tmp/new.txt' }, + }, + { + type: 'tool-approval-request', + approvalId: 'new-approval', + toolCallId: 'new-tool-call', + }, + { type: 'finish-step' }, + { type: 'finish' }, + ]); + + state = createStreamingUIMessageState({ + messageId: 'msg-123', + lastMessage: undefined, + }); + + await consumeStream({ + stream: processUIMessageStream({ + stream, + runUpdateMessageJob, + onError: error => { + throw error; + }, + }), + }); + + expect(state!.message.parts).toHaveLength(3); + expect(state!.message.parts.slice(0, 2)).toEqual([ + { type: 'step-start' }, + { type: 'step-start' }, + ]); + expect(state!.message.parts[2]).toMatchObject({ + type: 'tool-deleteFile', + toolCallId: 'new-tool-call', + state: 'approval-requested', + input: { path: '/tmp/new.txt' }, + approval: { id: 'new-approval' }, + }); + }); + describe('data ui parts (transient part)', () => { let dataCalls: InferUIMessageData[] = []; diff --git a/packages/ai/src/ui/process-ui-message-stream.ts b/packages/ai/src/ui/process-ui-message-stream.ts index 2b988d3a6c50..efe0fcb96110 100644 --- a/packages/ai/src/ui/process-ui-message-stream.ts +++ b/packages/ai/src/ui/process-ui-message-stream.ts @@ -104,7 +104,7 @@ export function processUIMessageStream({ new TransformStream>({ async transform(chunk, controller) { await runUpdateMessageJob(async ({ state, write }) => { - function getCurrentStepParts() { + function getCurrentStepStartIndex() { const parts = state.message.parts; let currentStepStartIndex = parts.length - 1; @@ -115,7 +115,12 @@ export function processUIMessageStream({ currentStepStartIndex--; } - return parts.slice(currentStepStartIndex + 1); + return currentStepStartIndex; + } + + function getCurrentStepParts() { + const currentStepStartIndex = getCurrentStepStartIndex(); + return state.message.parts.slice(currentStepStartIndex + 1); } function getCurrentStepToolInvocations() { @@ -921,6 +926,23 @@ export function processUIMessageStream({ default: { if (isDataUIMessageChunk(chunk)) { + if (chunk.type === 'data-reload') { + const currentStepStartIndex = getCurrentStepStartIndex(); + + // A retry invalidates everything streamed after the current + // step boundary. The retried step will add a new boundary + // before it starts writing output again. + state.message.parts = state.message.parts.slice( + 0, + currentStepStartIndex + 1, + ); + state.activeTextParts = createIdMap(); + state.activeReasoningParts = createIdMap(); + state.partialToolCalls = createIdMap(); + write(); + break; + } + // validate data chunk if dataPartSchemas is provided if (dataPartSchemas?.[chunk.type] != null) { const partIdx = state.message.parts.findIndex( diff --git a/packages/workflow/src/do-stream-step.test.ts b/packages/workflow/src/do-stream-step.test.ts new file mode 100644 index 000000000000..219dd5f948ee --- /dev/null +++ b/packages/workflow/src/do-stream-step.test.ts @@ -0,0 +1,79 @@ +import type { LanguageModelV4Prompt } from '@ai-sdk/provider'; +import type { Experimental_LanguageModelStreamPart, LanguageModel } from 'ai'; +import { convertArrayToReadableStream } from '@ai-sdk/provider-utils/test'; +import { describe, expect, it, vi } from 'vitest'; +import type * as AiModule from 'ai'; + +const { getStepMetadata, streamModelCall } = vi.hoisted(() => ({ + getStepMetadata: vi.fn(), + streamModelCall: vi.fn(), +})); + +vi.mock('workflow', () => ({ getStepMetadata })); +vi.mock('ai', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + experimental_streamLanguageModelCall: streamModelCall, + }; +}); + +const { doStreamStep } = await import('./do-stream-step.js'); + +describe('doStreamStep', () => { + it('does not emit an invalidation boundary for direct initial calls', async () => { + getStepMetadata.mockImplementationOnce(() => { + throw new Error( + '`getStepMetadata()` can only be called inside a step function', + ); + }); + streamModelCall.mockResolvedValueOnce({ + stream: convertArrayToReadableStream([]), + }); + + const chunks: unknown[] = []; + const writable = new WritableStream({ + write(chunk) { + chunks.push(chunk); + }, + }); + + await doStreamStep( + [] as unknown as LanguageModelV4Prompt, + {} as LanguageModel, + writable as WritableStream< + Experimental_LanguageModelStreamPart> + >, + ); + + expect(chunks).toEqual([]); + }); + + it('writes an invalidation boundary before a retried model stream', async () => { + getStepMetadata.mockReturnValue({ attempt: 2 }); + streamModelCall.mockResolvedValueOnce({ + stream: convertArrayToReadableStream([]), + }); + + const chunks: unknown[] = []; + const writable = new WritableStream({ + write(chunk) { + chunks.push(chunk); + }, + }); + + await doStreamStep( + [] as unknown as LanguageModelV4Prompt, + {} as LanguageModel, + writable as WritableStream< + Experimental_LanguageModelStreamPart> + >, + ); + + expect(chunks).toEqual([ + { type: 'finish-step' }, + { type: 'data-reload', data: {}, transient: true }, + { type: 'start-step' }, + ]); + }); +}); diff --git a/packages/workflow/src/do-stream-step.ts b/packages/workflow/src/do-stream-step.ts index b857c9504487..97ba73e9bd24 100644 --- a/packages/workflow/src/do-stream-step.ts +++ b/packages/workflow/src/do-stream-step.ts @@ -2,6 +2,7 @@ import type { LanguageModelV4CallOptions, LanguageModelV4Prompt, } from '@ai-sdk/provider'; +import { getStepMetadata } from 'workflow'; import { experimental_streamLanguageModelCall as streamModelCall, gateway, @@ -126,34 +127,6 @@ export async function doStreamStep( ? resolveSerializableTools(serializedTools) : undefined; - // streamModelCall handles: prompt standardization, tool preparation, - // model.doStream(), retry logic, and stream part transformation - // (tool call parsing, finish reason mapping, file wrapping). - const { stream: modelStream } = await streamModelCall({ - model, - // streamModelCall expects Prompt (ModelMessage[]) but we pass the - // pre-converted LanguageModelV4Prompt. standardizePrompt inside - // streamModelCall handles both formats. - messages: conversationPrompt as unknown as ModelMessage[], - allowSystemInMessages: true, - tools, - toolChoice: options?.toolChoice, - includeRawChunks: options?.includeRawChunks, - providerOptions: options?.providerOptions, - abortSignal: options?.abortSignal, - headers: options?.headers, - reasoning: options?.reasoning, - maxOutputTokens: options?.maxOutputTokens, - temperature: options?.temperature, - topP: options?.topP, - topK: options?.topK, - presencePenalty: options?.presencePenalty, - frequencyPenalty: options?.frequencyPenalty, - stopSequences: options?.stopSequences, - seed: options?.seed, - repairToolCall: options?.repairToolCall, - }); - // Consume the stream: capture data and write to writable in real-time const toolCalls: ParsedToolCall[] = []; const providerExecutedToolResults = new Map< @@ -174,6 +147,48 @@ export async function doStreamStep( const writer = writable?.getWriter(); try { + if (writer && isRetryAttempt()) { + await writer.write({ + type: 'finish-step', + } as unknown as ModelCallStreamPart); + await writer.write({ + type: 'data-reload', + data: {}, + transient: true, + } as unknown as ModelCallStreamPart); + await writer.write({ + type: 'start-step', + } as unknown as ModelCallStreamPart); + } + + // streamModelCall handles: prompt standardization, tool preparation, + // model.doStream(), retry logic, and stream part transformation + // (tool call parsing, finish reason mapping, file wrapping). + const { stream: modelStream } = await streamModelCall({ + model, + // streamModelCall expects Prompt (ModelMessage[]) but we pass the + // pre-converted LanguageModelV4Prompt. standardizePrompt inside + // streamModelCall handles both formats. + messages: conversationPrompt as unknown as ModelMessage[], + allowSystemInMessages: true, + tools, + toolChoice: options?.toolChoice, + includeRawChunks: options?.includeRawChunks, + providerOptions: options?.providerOptions, + abortSignal: options?.abortSignal, + headers: options?.headers, + reasoning: options?.reasoning, + maxOutputTokens: options?.maxOutputTokens, + temperature: options?.temperature, + topP: options?.topP, + topK: options?.topK, + presencePenalty: options?.presencePenalty, + frequencyPenalty: options?.frequencyPenalty, + stopSequences: options?.stopSequences, + seed: options?.seed, + repairToolCall: options?.repairToolCall, + }); + for await (const part of modelStream) { switch (part.type) { case 'text-delta': @@ -263,3 +278,21 @@ export async function doStreamStep( providerExecutedToolResults, }; } + +function isRetryAttempt() { + try { + return getStepMetadata().attempt > 1; + } catch (error) { + // Direct calls are supported by the compatibility tests and do not run + // inside a Workflow step, so they can only represent the first attempt. + if ( + error instanceof Error && + error.message === + '`getStepMetadata()` can only be called inside a step function' + ) { + return false; + } + + throw error; + } +} diff --git a/packages/workflow/src/to-ui-message-chunk.test.ts b/packages/workflow/src/to-ui-message-chunk.test.ts new file mode 100644 index 000000000000..8eded848f7f2 --- /dev/null +++ b/packages/workflow/src/to-ui-message-chunk.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from 'vitest'; +import { toUIMessageChunk } from './to-ui-message-chunk.js'; + +describe('toUIMessageChunk', () => { + it('preserves the retry invalidation marker as a transient data chunk', () => { + expect( + toUIMessageChunk({ + type: 'data-reload', + data: {}, + } as never), + ).toEqual({ + type: 'data-reload', + data: {}, + transient: true, + }); + }); +}); diff --git a/packages/workflow/src/to-ui-message-chunk.ts b/packages/workflow/src/to-ui-message-chunk.ts index b68ae8a7de64..833da4880951 100644 --- a/packages/workflow/src/to-ui-message-chunk.ts +++ b/packages/workflow/src/to-ui-message-chunk.ts @@ -199,6 +199,13 @@ export function toUIMessageChunk( toolCallId: passthroughPart.toolCallId, } as UIMessageChunk; } + if (passthroughPart.type === 'data-reload') { + return { + type: 'data-reload', + data: passthroughPart.data ?? {}, + transient: true, + } as UIMessageChunk; + } if ( passthroughPart.type === 'finish-step' || passthroughPart.type === 'start-step' || From e3c4d650a4f30eb188ac9f47a1707d88477293ad Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Mon, 10 Aug 2026 08:30:22 -0700 Subject: [PATCH 2/2] fix(workflow): isolate retry invalidation marker Signed-off-by: Ruiming Zhao --- .../ui-message-stream/ui-message-chunks.ts | 6 ++ .../src/ui/process-ui-message-stream.test.ts | 66 ++++++++++++++++++- .../ai/src/ui/process-ui-message-stream.ts | 40 ++++++----- packages/workflow/src/do-stream-step.test.ts | 2 +- packages/workflow/src/do-stream-step.ts | 4 +- .../workflow/src/to-ui-message-chunk.test.ts | 9 +-- packages/workflow/src/to-ui-message-chunk.ts | 8 +-- 7 files changed, 102 insertions(+), 33 deletions(-) diff --git a/packages/ai/src/ui-message-stream/ui-message-chunks.ts b/packages/ai/src/ui-message-stream/ui-message-chunks.ts index 2c590ecae74e..9ac9cbfaeed8 100644 --- a/packages/ai/src/ui-message-stream/ui-message-chunks.ts +++ b/packages/ai/src/ui-message-stream/ui-message-chunks.ts @@ -183,6 +183,9 @@ export const uiMessageChunkSchema = lazySchema(() => z.looseObject({ type: z.literal('finish-step'), }), + z.looseObject({ + type: z.literal('reload'), + }), z.looseObject({ type: z.literal('start'), messageId: z.string().optional(), @@ -378,6 +381,9 @@ export type UIMessageChunk< | { type: 'finish-step'; } + | { + type: 'reload'; + } | { type: 'start'; messageId?: string; diff --git a/packages/ai/src/ui/process-ui-message-stream.test.ts b/packages/ai/src/ui/process-ui-message-stream.test.ts index d964d0ddde24..89d3131645fd 100644 --- a/packages/ai/src/ui/process-ui-message-stream.test.ts +++ b/packages/ai/src/ui/process-ui-message-stream.test.ts @@ -4785,7 +4785,7 @@ describe('processUIMessageStream', () => { toolCallId: 'old-tool-call', }, { type: 'finish-step' }, - { type: 'data-reload', data: {}, transient: true }, + { type: 'reload' }, { type: 'start-step' }, { type: 'tool-input-available', @@ -4831,6 +4831,70 @@ describe('processUIMessageStream', () => { }); }); + it('should preserve a user data-reload part', async () => { + const stream = createUIMessageStream([ + { type: 'start', messageId: 'msg-123' }, + { type: 'start-step' }, + { + type: 'data-reload', + data: { source: 'application' }, + }, + { type: 'finish-step' }, + { type: 'finish' }, + ]); + + state = createStreamingUIMessageState({ + messageId: 'msg-123', + lastMessage: undefined, + }); + + await consumeStream({ + stream: processUIMessageStream({ + stream, + runUpdateMessageJob, + onError: error => { + throw error; + }, + }), + }); + + expect(state!.message.parts).toEqual([ + { type: 'step-start' }, + { type: 'data-reload', data: { source: 'application' } }, + ]); + }); + + it('should preserve existing parts when reload has no step boundary', async () => { + const stream = createUIMessageStream([ + { type: 'start', messageId: 'msg-123' }, + { type: 'reload' }, + { type: 'finish' }, + ]); + + state = createStreamingUIMessageState({ + messageId: 'msg-123', + lastMessage: { + role: 'assistant', + id: 'original-id', + parts: [{ type: 'text', text: 'existing response', state: 'done' }], + }, + }); + + await consumeStream({ + stream: processUIMessageStream({ + stream, + runUpdateMessageJob, + onError: error => { + throw error; + }, + }), + }); + + expect(state!.message.parts).toEqual([ + { type: 'text', text: 'existing response', state: 'done' }, + ]); + }); + describe('data ui parts (transient part)', () => { let dataCalls: InferUIMessageData[] = []; diff --git a/packages/ai/src/ui/process-ui-message-stream.ts b/packages/ai/src/ui/process-ui-message-stream.ts index efe0fcb96110..1819672879bd 100644 --- a/packages/ai/src/ui/process-ui-message-stream.ts +++ b/packages/ai/src/ui/process-ui-message-stream.ts @@ -127,6 +127,23 @@ export function processUIMessageStream({ return getCurrentStepParts().filter(isToolUIPart); } + function resetCurrentStep() { + const currentStepStartIndex = getCurrentStepStartIndex(); + + // A retry invalidates everything streamed after the current step + // boundary. If no boundary exists, keep the message intact. + if (currentStepStartIndex >= 0) { + state.message.parts = state.message.parts.slice( + 0, + currentStepStartIndex + 1, + ); + } + + state.activeTextParts = createIdMap(); + state.activeReasoningParts = createIdMap(); + state.partialToolCalls = createIdMap(); + } + function getToolInvocation(toolCallId: string) { const toolInvocations = getCurrentStepToolInvocations(); @@ -924,25 +941,14 @@ export function processUIMessageStream({ break; } + case 'reload': { + resetCurrentStep(); + write(); + break; + } + default: { if (isDataUIMessageChunk(chunk)) { - if (chunk.type === 'data-reload') { - const currentStepStartIndex = getCurrentStepStartIndex(); - - // A retry invalidates everything streamed after the current - // step boundary. The retried step will add a new boundary - // before it starts writing output again. - state.message.parts = state.message.parts.slice( - 0, - currentStepStartIndex + 1, - ); - state.activeTextParts = createIdMap(); - state.activeReasoningParts = createIdMap(); - state.partialToolCalls = createIdMap(); - write(); - break; - } - // validate data chunk if dataPartSchemas is provided if (dataPartSchemas?.[chunk.type] != null) { const partIdx = state.message.parts.findIndex( diff --git a/packages/workflow/src/do-stream-step.test.ts b/packages/workflow/src/do-stream-step.test.ts index 219dd5f948ee..bc37bc6160b4 100644 --- a/packages/workflow/src/do-stream-step.test.ts +++ b/packages/workflow/src/do-stream-step.test.ts @@ -72,7 +72,7 @@ describe('doStreamStep', () => { expect(chunks).toEqual([ { type: 'finish-step' }, - { type: 'data-reload', data: {}, transient: true }, + { type: 'reload' }, { type: 'start-step' }, ]); }); diff --git a/packages/workflow/src/do-stream-step.ts b/packages/workflow/src/do-stream-step.ts index 97ba73e9bd24..5a9e967560c3 100644 --- a/packages/workflow/src/do-stream-step.ts +++ b/packages/workflow/src/do-stream-step.ts @@ -152,9 +152,7 @@ export async function doStreamStep( type: 'finish-step', } as unknown as ModelCallStreamPart); await writer.write({ - type: 'data-reload', - data: {}, - transient: true, + type: 'reload', } as unknown as ModelCallStreamPart); await writer.write({ type: 'start-step', diff --git a/packages/workflow/src/to-ui-message-chunk.test.ts b/packages/workflow/src/to-ui-message-chunk.test.ts index 8eded848f7f2..1395126dc416 100644 --- a/packages/workflow/src/to-ui-message-chunk.test.ts +++ b/packages/workflow/src/to-ui-message-chunk.test.ts @@ -2,16 +2,13 @@ import { describe, expect, it } from 'vitest'; import { toUIMessageChunk } from './to-ui-message-chunk.js'; describe('toUIMessageChunk', () => { - it('preserves the retry invalidation marker as a transient data chunk', () => { + it('preserves the retry invalidation marker as a UI chunk', () => { expect( toUIMessageChunk({ - type: 'data-reload', - data: {}, + type: 'reload', } as never), ).toEqual({ - type: 'data-reload', - data: {}, - transient: true, + type: 'reload', }); }); }); diff --git a/packages/workflow/src/to-ui-message-chunk.ts b/packages/workflow/src/to-ui-message-chunk.ts index 833da4880951..785c9fcad311 100644 --- a/packages/workflow/src/to-ui-message-chunk.ts +++ b/packages/workflow/src/to-ui-message-chunk.ts @@ -199,12 +199,10 @@ export function toUIMessageChunk( toolCallId: passthroughPart.toolCallId, } as UIMessageChunk; } - if (passthroughPart.type === 'data-reload') { + if (passthroughPart.type === 'reload') { return { - type: 'data-reload', - data: passthroughPart.data ?? {}, - transient: true, - } as UIMessageChunk; + type: 'reload', + } as unknown as UIMessageChunk; } if ( passthroughPart.type === 'finish-step' ||