fix(workflow): invalidate UI parts when a step retries

.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.

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;

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<UIMessage>[] = [];
 

packages/ai/src/ui/process-ui-message-stream.ts

@@ -104,7 +104,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
     new TransformStream<UIMessageChunk, InferUIMessageChunk<UI_MESSAGE>>({
       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<UI_MESSAGE extends UIMessage>({
               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<UI_MESSAGE extends UIMessage>({
               break;
             }
 
+            case 'reload': {
+              resetCurrentStep();
+              write();
+              break;
+            }
+
             default: {
               if (isDataUIMessageChunk(chunk)) {
                 // validate data chunk if dataPartSchemas is provided

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<typeof AiModule>();
+  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<Record<string, never>>
+      >,
+    );
+
+    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<Record<string, never>>
+      >,
+    );
+
+    expect(chunks).toEqual([
+      { type: 'finish-step' },
+      { type: 'reload' },
+      { type: 'start-step' },
+    ]);
+  });
+});