feat(ai): add onStepFinish to agent.generate and agent.stream (#11980)

## Background

For things such as token tracking, it would be helpful if callbacks can
be registered on a per-call basis. See #11468

## Summary

Add `onStepFinish` callback support to `Agent.generate()` and
`Agent.stream()`

## Related Issues

Resolves #11468

Author: lgrammel

.changeset/stale-schools-behave.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat(ai): add onStepFinish to agent.generate and agent.stream

content/docs/03-agents/02-building-agents.mdx

@@ -329,6 +329,45 @@ export async function POST(request: Request) {
 }
 ```
 
+### Track Step Progress
+
+Use `onStepFinish` to track each step's progress, including token usage:
+
+```ts
+const result = await myAgent.generate({
+  prompt: 'Research and summarize the latest AI trends',
+  onStepFinish: async ({ usage, finishReason, toolCalls }) => {
+    console.log('Step completed:', {
+      inputTokens: usage.inputTokens,
+      outputTokens: usage.outputTokens,
+      finishReason,
+      toolsUsed: toolCalls?.map(tc => tc.toolName),
+    });
+  },
+});
+```
+
+You can also define `onStepFinish` in the constructor for agent-wide tracking. When both constructor and method callbacks are provided, both are called (constructor first, then the method callback):
+
+```ts
+const agent = new ToolLoopAgent({
+  model: __MODEL__,
+  onStepFinish: async ({ usage }) => {
+    // Agent-wide logging
+    console.log('Agent step:', usage.totalTokens);
+  },
+});
+
+// Method-level callback runs after constructor callback
+const result = await agent.generate({
+  prompt: 'Hello',
+  onStepFinish: async ({ usage }) => {
+    // Per-call tracking (e.g., for billing)
+    await trackUsage(usage);
+  },
+});
+```
+
 ## End-to-end Type Safety
 
 You can infer types for your agent's `UIMessage`s:

content/docs/07-reference/01-ai-sdk-core/15-agent.mdx

@@ -18,7 +18,9 @@ import { Output } from '../generate-text/output';
 import { GenerateTextResult } from '../generate-text/generate-text-result';
 import { StreamTextResult } from '../generate-text/stream-text-result';
 
-export type AgentCallParameters<CALL_OPTIONS> = ([CALL_OPTIONS] extends [never]
+export type AgentCallParameters<CALL_OPTIONS, TOOLS extends ToolSet = {}> = ([
+  CALL_OPTIONS,
+] extends [never]
   ? { options?: never }
   : { options: CALL_OPTIONS }) &
   (
@@ -63,6 +65,10 @@ export type AgentCallParameters<CALL_OPTIONS> = ([CALL_OPTIONS] extends [never]
      * Can be used alongside abortSignal.
      */
     timeout?: number | { totalMs?: number };
+    /**
+     * Callback that is called when each step (LLM call) is finished, including intermediate steps.
+     */
+    onStepFinish?: ToolLoopAgentOnStepFinishCallback<TOOLS>;
   };
 
 /**
@@ -97,14 +103,14 @@ export interface Agent<
    * Generates an output from the agent (non-streaming).
    */
   generate(
-    options: AgentCallParameters<CALL_OPTIONS>,
+    options: AgentCallParameters<CALL_OPTIONS, TOOLS>,
   ): PromiseLike<GenerateTextResult<TOOLS, OUTPUT>>;
 
   /**
    * Streams an output from the agent (streaming).
    */
   stream(
-    options: AgentCallParameters<CALL_OPTIONS>,
+    options: AgentStreamParameters<CALL_OPTIONS, TOOLS>,
   ): PromiseLike<StreamTextResult<TOOLS, OUTPUT>>;
 }
 ```
@@ -129,13 +135,14 @@ export interface Agent<
 
 ## Method Parameters
 
-Both `generate()` and `stream()` accept an `AgentCallParameters<CALL_OPTIONS>` object with:
+Both `generate()` and `stream()` accept an `AgentCallParameters<CALL_OPTIONS, TOOLS>` object with:
 
 - `prompt` (optional): A string prompt or array of `ModelMessage` objects
 - `messages` (optional): An array of `ModelMessage` objects (mutually exclusive with `prompt`)
 - `options` (optional): Additional call options when `CALL_OPTIONS` is not `never`
 - `abortSignal` (optional): An `AbortSignal` to cancel the operation
 - `timeout` (optional): A timeout in milliseconds. Can be specified as a number or as an object with a `totalMs` property. The call will be aborted if it takes longer than the specified timeout. Can be used alongside `abortSignal`.
+- `onStepFinish` (optional): A callback invoked after each agent step (LLM/tool call) completes. Useful for tracking token usage or logging.
 
 ## Example: Custom Agent Implementation
 

content/docs/07-reference/01-ai-sdk-core/16-tool-loop-agent.mdx

@@ -248,6 +248,13 @@ const result = await agent.generate({
       description:
         'Timeout in milliseconds. Can be specified as a number or as an object with a totalMs property. The call will be aborted if it takes longer than the specified timeout. Can be used alongside abortSignal.',
     },
+    {
+      name: 'onStepFinish',
+      type: 'ToolLoopAgentOnStepFinishCallback',
+      isOptional: true,
+      description:
+        'Callback invoked after each agent step (LLM/tool call) completes. If also specified in the constructor, both callbacks are called (constructor first, then this one).',
+    },
   ]}
 />
 
@@ -302,6 +309,13 @@ for await (const chunk of stream.textStream) {
       description:
         'Optional stream transformation(s). They are applied in the order provided and must maintain the stream structure. See `streamText` docs for details.',
     },
+    {
+      name: 'onStepFinish',
+      type: 'ToolLoopAgentOnStepFinishCallback',
+      isOptional: true,
+      description:
+        'Callback invoked after each agent step (LLM/tool call) completes. If also specified in the constructor, both callbacks are called (constructor first, then this one).',
+    },
   ]}
 />
 

content/docs/07-reference/01-ai-sdk-core/17-create-agent-ui-stream.mdx

@@ -86,6 +86,13 @@ export async function* streamAgent(
       description:
         'Optional transformations to apply to the agent output stream (experimental).',
     },
+    {
+      name: 'onStepFinish',
+      type: 'ToolLoopAgentOnStepFinishCallback',
+      isRequired: false,
+      description:
+        'Callback invoked after each agent step (LLM/tool call) completes. Useful for tracking token usage or logging intermediate steps.',
+    },
     {
       name: '...UIMessageStreamOptions',
       type: 'UIMessageStreamOptions',

content/docs/07-reference/01-ai-sdk-core/18-create-agent-ui-stream-response.mdx

@@ -87,6 +87,13 @@ export async function POST(request: Request) {
       description:
         'Optional stream transforms to post-process text output—the same as in lower-level streaming APIs.',
     },
+    {
+      name: 'onStepFinish',
+      type: 'ToolLoopAgentOnStepFinishCallback',
+      isRequired: false,
+      description:
+        'Callback invoked after each agent step (LLM/tool call) completes. Useful for tracking token usage or logging intermediate steps.',
+    },
     {
       name: '...UIMessageStreamOptions',
       type: 'UIMessageStreamOptions',

content/docs/07-reference/01-ai-sdk-core/18-pipe-agent-ui-stream-to-response.mdx

@@ -88,6 +88,13 @@ export async function handler(req, res) {
       description:
         'Optional stream text transformation(s) applied to agent output.',
     },
+    {
+      name: 'onStepFinish',
+      type: 'ToolLoopAgentOnStepFinishCallback',
+      isRequired: false,
+      description:
+        'Callback invoked after each agent step (LLM/tool call) completes. Useful for tracking token usage or logging intermediate steps.',
+    },
     {
       name: '...UIMessageStreamResponseInit & UIMessageStreamOptions',
       type: 'object',

examples/ai-functions/src/agent/openai-generate-on-step-finish.ts

@@ -0,0 +1,55 @@
+import { openai } from '@ai-sdk/openai';
+import { tool, ToolLoopAgent } from 'ai';
+import { z } from 'zod';
+import { run } from '../lib/run';
+
+const agent = new ToolLoopAgent({
+  model: openai('gpt-4o'),
+  instructions: 'You are a helpful assistant that can look up weather.',
+  tools: {
+    weather: tool({
+      description: 'Get the weather for a location',
+      inputSchema: z.object({
+        location: z.string().describe('The location to get weather for'),
+      }),
+      execute: ({ location }) => ({
+        location,
+        temperature: 72 + Math.floor(Math.random() * 21) - 10,
+        condition: 'sunny',
+      }),
+    }),
+  },
+});
+
+run(async () => {
+  const result = await agent.generate({
+    prompt: 'What is the weather in San Francisco and New York?',
+    onStepFinish: async ({ text, finishReason, usage, toolCalls }) => {
+      console.log('\n--- Step Finished ---');
+      console.log('Finish Reason:', finishReason);
+      console.log('Token Usage:', {
+        input: usage.inputTokens,
+        output: usage.outputTokens,
+        total: usage.totalTokens,
+      });
+
+      if (toolCalls && toolCalls.length > 0) {
+        console.log(
+          'Tool Calls:',
+          toolCalls.map(tc => tc.toolName),
+        );
+      }
+
+      if (text) {
+        console.log(
+          'Text:',
+          text.substring(0, 100) + (text.length > 100 ? '...' : ''),
+        );
+      }
+    },
+  });
+
+  console.log('\n=== Final Result ===');
+  console.log('Total Steps:', result.steps.length);
+  console.log('Final Text:', result.text);
+});

packages/ai/src/agent/agent.ts

@@ -5,11 +5,14 @@ import { StreamTextTransform } from '../generate-text/stream-text';
 import { StreamTextResult } from '../generate-text/stream-text-result';
 import { ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
+import { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-on-step-finish-callback';
 
 /**
  * Parameters for calling an agent.
  */
-export type AgentCallParameters<CALL_OPTIONS> = ([CALL_OPTIONS] extends [never]
+export type AgentCallParameters<CALL_OPTIONS, TOOLS extends ToolSet = {}> = ([
+  CALL_OPTIONS,
+] extends [never]
   ? { options?: never }
   : { options: CALL_OPTIONS }) &
   (
@@ -53,6 +56,11 @@ export type AgentCallParameters<CALL_OPTIONS> = ([CALL_OPTIONS] extends [never]
      * Timeout in milliseconds. Can be specified as a number or as an object with `totalMs`.
      */
     timeout?: TimeoutConfiguration;
+
+    /**
+     * Callback that is called when each step (LLM call) is finished, including intermediate steps.
+     */
+    onStepFinish?: ToolLoopAgentOnStepFinishCallback<TOOLS>;
   };
 
 /**
@@ -61,7 +69,7 @@ export type AgentCallParameters<CALL_OPTIONS> = ([CALL_OPTIONS] extends [never]
 export type AgentStreamParameters<
   CALL_OPTIONS,
   TOOLS extends ToolSet,
-> = AgentCallParameters<CALL_OPTIONS> & {
+> = AgentCallParameters<CALL_OPTIONS, TOOLS> & {
   /**
    * Optional stream transformations.
    * They are applied in the order they are provided.
@@ -104,7 +112,7 @@ export interface Agent<
    * Generates an output from the agent (non-streaming).
    */
   generate(
-    options: AgentCallParameters<CALL_OPTIONS>,
+    options: AgentCallParameters<CALL_OPTIONS, TOOLS>,
   ): PromiseLike<GenerateTextResult<TOOLS, OUTPUT>>;
 
   /**

packages/ai/src/agent/create-agent-ui-stream-response.ts

@@ -7,12 +7,14 @@ import { UIMessageStreamResponseInit } from '../ui-message-stream/ui-message-str
 import { InferUITools, UIMessage } from '../ui/ui-messages';
 import { Agent } from './agent';
 import { createAgentUIStream } from './create-agent-ui-stream';
+import { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-on-step-finish-callback';
 
 /**
  * Runs the agent and returns a response object with a UI message stream.
  *
  * @param agent - The agent to run.
  * @param uiMessages - The input UI messages.
+ * @param onStepFinish - Callback that is called when each step is finished. Optional.
  *
  * @returns The response object.
  */
@@ -36,6 +38,7 @@ export async function createAgentUIStreamResponse<
   experimental_transform?:
     | StreamTextTransform<TOOLS>
     | Array<StreamTextTransform<TOOLS>>;
+  onStepFinish?: ToolLoopAgentOnStepFinishCallback<TOOLS>;
 } & UIMessageStreamResponseInit &
   UIMessageStreamOptions<
     UIMessage<MESSAGE_METADATA, never, InferUITools<TOOLS>>