feat(ai): add onStepFinish continuation support for validation and retry

.changeset/spotty-impalas-whisper.md

@@ -0,0 +1,9 @@
+---
+'ai': patch
+---
+
+feat(ai): add onStepFinish continuation support for validation and retry
+
+Add ability for `onStepFinish` callback to return `StepContinueResult` to continue the generation loop with injected feedback messages. This enables validation and automatic retry functionality for `generateText`, `streamText`, and `generateObject`.
+
+The callback signature is backward compatible - existing code returning `void` continues to work unchanged.

content/cookbook/05-node/47-generate-object-continuation.mdx

@@ -0,0 +1,80 @@
+---
+title: Generate Object with Validation and Retries
+description: Learn how to use generateObject with onStepFinish for validation and automatic retries.
+tags: ['node', 'structured data', 'validation']
+---
+
+# Generate Object with Validation and Retries
+
+You can use the `onStepFinish` callback in `generateObject` to validate the generated object and continue the generation loop with feedback if validation fails. This allows you to implement self-correcting loops where the model can fix its mistakes based on your specific validation logic.
+
+This is particularly useful when:
+
+1. You have complex validation rules that cannot be expressed in the schema alone (e.g., "age must be between 18 and 120").
+2. You want to give the model specific feedback on _why_ the validation failed.
+3. You want to automatically retry generation a limited number of times.
+
+## Example
+
+In this example, we'll generate a user object and validate that the name meets length requirements and the age is within a valid range.
+
+```ts file='index.ts'
+import { generateObject, type StepContinueResult } from 'ai';
+import { openai } from '@ai-sdk/openai';
+import { z } from 'zod';
+
+const result = await generateObject({
+  model: openai('gpt-4o'),
+  schema: z.object({
+    name: z.string(),
+    age: z.number(),
+    email: z.string().email(),
+  }),
+  prompt: 'Generate a user object for a test user.',
+  maxRetries: 5, // Safety limit: max 5 attempts
+  onStepFinish: async (step): Promise<StepContinueResult> => {
+    // If the schema validation failed (e.g. malformed JSON or type mismatch),
+    // step.validationError will be populated.
+    // You can also perform additional custom validation on step.object.
+
+    const issues: string[] = [];
+
+    // Check for schema validation errors
+    if (step.validationError) {
+      issues.push(`Schema validation failed: ${step.validationError.message}`);
+    }
+
+    // Custom validation logic on the successfully parsed object
+    if (step.object) {
+      if (step.object.name.length < 3 || step.object.name.length > 50) {
+        issues.push('Name must be between 3 and 50 characters');
+      }
+      if (step.object.age < 18 || step.object.age > 120) {
+        issues.push('Age must be between 18 and 120');
+      }
+    }
+
+    // If there are issues, continue with feedback
+    if (issues.length > 0) {
+      console.log(
+        `Validation failed, retrying... Issues: ${issues.join(', ')}`,
+      );
+      return {
+        continue: true,
+        messages: [
+          {
+            role: 'user',
+            content: `Please fix the following issues: ${issues.join(', ')}`,
+          },
+        ],
+      };
+    }
+
+    console.log('Validation passed!');
+    // If validation succeeded, return { continue: false } to stop the loop
+    return { continue: false };
+  },
+});
+
+console.log(JSON.stringify(result.object, null, 2));
+```

content/docs/03-ai-sdk-core/05-generating-text.mdx

@@ -234,6 +234,50 @@ const result = streamText({
 });
 ```
 
+### `onStepFinish` callback
+
+When using `generateText` or `streamText`, you can provide an `onStepFinish` callback that is triggered when a step is finished.
+It contains all the text, tool calls, and tool results for the step.
+When you have multiple steps, the callback is triggered for each step.
+
+#### Continuing the loop with feedback
+
+You can return a `StepContinueResult` from `onStepFinish` to continue the generation loop with injected feedback messages.
+This is useful for validating outputs and automatically retrying when validation fails.
+
+```tsx highlight="5-20"
+import { generateText, stepCountIs, type StepContinueResult } from 'ai';
+
+const result = await generateText({
+  model: 'openai/gpt-4o-mini',
+  prompt: 'Generate a text message for SMS: no markdown, under 160 characters',
+  onStepFinish: async (step): Promise<StepContinueResult> => {
+    const text = step.text;
+    const hasMarkdown = /[*_`\[\]#]/.test(text);
+    const tooLong = text.length > 160;
+
+    if (hasMarkdown || tooLong) {
+      return {
+        continue: true,
+        messages: [
+          {
+            role: 'user',
+            content: `Validation failed: The message ${
+              hasMarkdown ? 'contains markdown symbols' : ''
+            }${hasMarkdown && tooLong ? ' and ' : ''}${
+              tooLong ? `is ${text.length} characters (max 160)` : ''
+            }. Please regenerate a plain text message without markdown and under 160 characters.`,
+          },
+        ],
+      };
+    }
+
+    return { continue: false };
+  },
+  stopWhen: stepCountIs(5), // Safety limit: max 5 attempts
+});
+```
+
 ### `fullStream` property
 
 You can read a stream with all events using the `fullStream` property.

content/docs/03-ai-sdk-core/10-generating-structured-data.mdx

@@ -274,6 +274,43 @@ try {
 }
 ```
 
+## Continuation with Validation
+
+You can use the `onStepFinish` callback to validate the generated object and continue the generation loop with feedback if validation fails.
+This is useful for cases where the model might generate an object that is valid JSON but does not meet your specific requirements (e.g. logical constraints).
+
+```ts
+import { generateObject, type StepContinueResult } from 'ai';
+import { z } from 'zod';
+
+const result = await generateObject({
+  model: 'openai/gpt-4o-mini',
+  schema: z.object({
+    name: z.string().min(3).max(50),
+    age: z.number().int().min(18).max(120),
+  }),
+  prompt: 'Generate a user object with name (3-50 chars) and age (18-120)',
+  onStepFinish: async (step): Promise<StepContinueResult> => {
+    // if validation failed, you can return feedback to the model
+    if (step.validationError) {
+      return {
+        continue: true,
+        messages: [
+          {
+            role: 'user',
+            content: `Validation failed: ${step.validationError.message}. Please regenerate the object with valid values.`,
+          },
+        ],
+      };
+    }
+
+    // if validation succeeded, stop the loop
+    return { continue: false };
+  },
+  maxRetries: 5, // Safety limit: max 5 attempts
+});
+```
+
 ## Repairing Invalid or Malformed JSON
 
 <Note type="warning">

content/docs/03-ai-sdk-core/15-tools-and-tool-calling.mdx

@@ -140,6 +140,49 @@ const result = await generateText({
 });
 ```
 
+#### Continuing the loop with feedback
+
+You can return a `StepContinueResult` from `onStepFinish` to continue the generation loop with injected feedback messages. This is useful for validating outputs and automatically retrying when validation fails.
+
+```tsx highlight="5-20"
+import { generateText, stepCountIs, type StepContinueResult } from 'ai';
+
+const result = await generateText({
+  model: 'openai/gpt-4o-mini',
+  prompt: 'Generate a text message for SMS: no markdown, under 160 characters',
+  onStepFinish: async (step): Promise<StepContinueResult> => {
+    const text = step.text;
+    const hasMarkdown = /[*_`\[\]#]/.test(text);
+    const tooLong = text.length > 160;
+
+    if (hasMarkdown || tooLong) {
+      return {
+        continue: true,
+        messages: [
+          {
+            role: 'user',
+            content: `Validation failed: The message ${
+              hasMarkdown ? 'contains markdown symbols' : ''
+            }${hasMarkdown && tooLong ? ' and ' : ''}${
+              tooLong ? `is ${text.length} characters (max 160)` : ''
+            }. Please regenerate a plain text message without markdown and under 160 characters.`,
+          },
+        ],
+      };
+    }
+
+    return { continue: false };
+  },
+  stopWhen: stepCountIs(5), // Safety limit: max 5 attempts
+});
+```
+
+The `StepContinueResult` type can be:
+
+- `{ continue: true, messages: Array<ModelMessage> }` - Continue the loop with injected feedback messages
+- `{ continue: false }` - Stop the loop (even if tool calls exist)
+- `undefined` or `void` - Continue normally based on tool calls and stop conditions
+
 ### `prepareStep` callback
 
 The `prepareStep` callback is called before a step is started.

content/docs/07-reference/01-ai-sdk-core/01-generate-text.mdx

@@ -739,9 +739,10 @@ To see `generateText` in action, check out [these examples](#examples).
     },
     {
       name: 'onStepFinish',
-      type: '(result: OnStepFinishResult) => Promise<void> | void',
+      type: '(result: OnStepFinishResult) => Promise<StepContinueResult> | StepContinueResult | Promise<void> | void',
       isOptional: true,
-      description: 'Callback that is called when a step is finished.',
+      description:
+        'Callback that is called when a step is finished. Optionally returns a `StepContinueResult` to continue the loop with feedback messages. If `void` or `undefined` is returned, the loop continues normally based on tool calls and stop conditions. If `{ continue: true, messages }` is returned, the loop continues with the injected messages. If `{ continue: false }` is returned, the loop stops (even if tool calls exist).',
       properties: [
         {
           type: 'OnStepFinishResult',

content/docs/07-reference/01-ai-sdk-core/02-stream-text.mdx

@@ -989,9 +989,10 @@ To see `streamText` in action, check out [these examples](#examples).
     },
     {
       name: 'onStepFinish',
-      type: '(result: onStepFinishResult) => Promise<void> | void',
+      type: '(result: onStepFinishResult) => Promise<StepContinueResult> | StepContinueResult | Promise<void> | void',
       isOptional: true,
-      description: 'Callback that is called when a step is finished.',
+      description:
+        'Callback that is called when a step is finished. Optionally returns a `StepContinueResult` to continue the loop with feedback messages. If `void` or `undefined` is returned, the loop continues normally based on tool calls and stop conditions. If `{ continue: true, messages }` is returned, the loop continues with the injected messages. If `{ continue: false }` is returned, the loop stops (even if tool calls exist).',
       properties: [
         {
           type: 'onStepFinishResult',

content/docs/07-reference/01-ai-sdk-core/03-generate-object.mdx

@@ -567,6 +567,74 @@ To see `generateObject` in action, check out the [additional examples](#more-exa
       description:
         'Provider-specific options. The outer key is the provider name. The inner values are the metadata. Details depend on the provider.',
     },
+    {
+      name: 'onStepFinish',
+      type: '(result: OnStepFinishResult) => Promise<StepContinueResult> | StepContinueResult | Promise<void> | void',
+      isOptional: true,
+      description:
+        'Callback that is called when a step is finished. Optionally returns a `StepContinueResult` to continue the loop with feedback messages. If `void` or `undefined` is returned, the loop continues normally based on tool calls and stop conditions. If `{ continue: true, messages }` is returned, the loop continues with the injected messages. If `{ continue: false }` is returned, the loop stops.',
+      properties: [
+        {
+          type: 'OnStepFinishResult',
+          parameters: [
+            {
+              name: 'object',
+              type: 'RESULT | undefined',
+              description:
+                'The generated object, if validation was successful. undefined if validation failed.',
+            },
+            {
+              name: 'validationError',
+              type: 'Error | undefined',
+              description:
+                'The validation error, if validation failed. undefined if validation was successful.',
+            },
+            {
+              name: 'text',
+              type: 'string',
+              description: 'The raw text generated by the model.',
+            },
+            {
+              name: 'finishReason',
+              type: '"stop" | "length" | "content-filter" | "tool-calls" | "error" | "other" | "unknown"',
+              description:
+                'The reason the model finished generating the text for the step.',
+            },
+            {
+              name: 'usage',
+              type: 'LanguageModelUsage',
+              description: 'The token usage of the step.',
+            },
+            {
+              name: 'warnings',
+              type: 'CallWarning[] | undefined',
+              description: 'Warnings from the model provider for this step.',
+            },
+            {
+              name: 'request',
+              type: 'LanguageModelRequestMetadata',
+              description: 'Request metadata for this step.',
+            },
+            {
+              name: 'response',
+              type: 'LanguageModelResponseMetadata',
+              description: 'Response metadata for this step.',
+            },
+            {
+              name: 'providerMetadata',
+              type: 'ProviderMetadata | undefined',
+              description:
+                'Additional provider-specific metadata for this step.',
+            },
+            {
+              name: 'reasoning',
+              type: 'string | undefined',
+              description: 'The reasoning text for this step.',
+            },
+          ],
+        },
+      ],
+    },
   ]}
 />