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-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 b9e537d30310..89d3131645fd 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,132 @@ 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: 'reload' }, + { 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' }, + }); + }); + + 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 2b988d3a6c50..1819672879bd 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,13 +115,35 @@ 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() { 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(); @@ -919,6 +941,12 @@ export function processUIMessageStream({ break; } + case 'reload': { + resetCurrentStep(); + write(); + break; + } + default: { if (isDataUIMessageChunk(chunk)) { // validate data chunk if dataPartSchemas is provided 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..bc37bc6160b4 --- /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: 'reload' }, + { type: 'start-step' }, + ]); + }); +}); diff --git a/packages/workflow/src/do-stream-step.ts b/packages/workflow/src/do-stream-step.ts index b94af38eb0c9..3180d9ab6526 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, @@ -147,36 +148,6 @@ export async function doStreamStep( return 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, - output, - 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< @@ -197,6 +168,47 @@ 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: 'reload', + } 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, + output, + 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': @@ -286,3 +298,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..1395126dc416 --- /dev/null +++ b/packages/workflow/src/to-ui-message-chunk.test.ts @@ -0,0 +1,14 @@ +import { describe, expect, it } from 'vitest'; +import { toUIMessageChunk } from './to-ui-message-chunk.js'; + +describe('toUIMessageChunk', () => { + it('preserves the retry invalidation marker as a UI chunk', () => { + expect( + toUIMessageChunk({ + type: 'reload', + } as never), + ).toEqual({ + type: 'reload', + }); + }); +}); diff --git a/packages/workflow/src/to-ui-message-chunk.ts b/packages/workflow/src/to-ui-message-chunk.ts index b68ae8a7de64..785c9fcad311 100644 --- a/packages/workflow/src/to-ui-message-chunk.ts +++ b/packages/workflow/src/to-ui-message-chunk.ts @@ -199,6 +199,11 @@ export function toUIMessageChunk( toolCallId: passthroughPart.toolCallId, } as UIMessageChunk; } + if (passthroughPart.type === 'reload') { + return { + type: 'reload', + } as unknown as UIMessageChunk; + } if ( passthroughPart.type === 'finish-step' || passthroughPart.type === 'start-step' ||