feat(core): add prepareStep to AgentOptions for per-step tool control (#1192)

Surfaces the AI SDK's prepareStep callback as a top-level AgentOptions
property. Users can now set a default step preparation callback at
agent creation time to control tool availability per step.

Per-call prepareStep in method options overrides the agent-level default.
The agent-level default is applied before applyForcedToolChoice so both
features compose correctly.

Fixes #1187

Author: ravyg

.changeset/feat-agent-prepare-step.md

@@ -0,0 +1,9 @@
+---
+"@voltagent/core": minor
+---
+
+feat(core): add `prepareStep` to AgentOptions for per-step tool control
+
+Surfaces the AI SDK's `prepareStep` callback as a top-level `AgentOptions` property so users can set a default step preparation callback at agent creation time. Per-call `prepareStep` in method options overrides the agent-level default.
+
+This enables controlling tool availability, tool choice, and other step settings on a per-step basis without passing `prepareStep` on every call.

packages/core/src/agent/agent.ts

@@ -128,7 +128,7 @@ export type {
   SemanticMemoryOptions,
 } from "./types";
 import { P, match } from "ts-pattern";
-import type { StopWhen } from "../ai-types";
+import type { PrepareStep, StopWhen } from "../ai-types";
 import type { SamplingPolicy } from "../eval/runtime";
 import type { ConversationStepRecord } from "../memory/types";
 import { applySummarization } from "./apply-summarization";
@@ -923,6 +923,13 @@ export interface BaseGenerationOptions<TProviderOptions extends ProviderOptions
    * Tool choice strategy for AI SDK calls.
    */
   toolChoice?: ToolChoice<Record<string, unknown>>;
+
+  /**
+   * Step preparation callback (ai-sdk `prepareStep`).
+   * Called before each step to control tool availability, tool choice, etc.
+   * Overrides the agent-level `prepareStep` if provided.
+   */
+  prepareStep?: PrepareStep;
 }
 
 export type GenerateTextOptions<
@@ -964,6 +971,7 @@ export class Agent {
   readonly maxSteps: number;
   readonly maxRetries: number;
   readonly stopWhen?: StopWhen;
+  readonly prepareStep?: PrepareStep;
   readonly markdown: boolean;
   readonly inheritParentSpan: boolean;
   readonly voice?: Voice;
@@ -1022,6 +1030,7 @@ export class Agent {
     this.maxSteps = options.maxSteps ?? defaultMaxSteps;
     this.maxRetries = options.maxRetries ?? DEFAULT_LLM_MAX_RETRIES;
     this.stopWhen = options.stopWhen;
+    this.prepareStep = options.prepareStep;
     this.markdown = options.markdown ?? false;
     this.inheritParentSpan = options.inheritParentSpan ?? true;
     this.voice = options.voice;
@@ -1265,6 +1274,11 @@ export class Agent {
               ...aiSDKOptions
             } = options || {};
 
+            // Apply agent-level prepareStep as default (per-call overrides)
+            if (this.prepareStep && !aiSDKOptions.prepareStep) {
+              aiSDKOptions.prepareStep = this.prepareStep as AITextCallOptions["prepareStep"];
+            }
+
             const forcedToolChoice = oc.systemContext.get(FORCED_TOOL_CHOICE_CONTEXT_KEY) as
               | ToolChoice<Record<string, unknown>>
               | undefined;
@@ -1879,6 +1893,11 @@ export class Agent {
           ...aiSDKOptions
         } = options || {};
 
+        // Apply agent-level prepareStep as default (per-call overrides)
+        if (this.prepareStep && !aiSDKOptions.prepareStep) {
+          aiSDKOptions.prepareStep = this.prepareStep as AITextCallOptions["prepareStep"];
+        }
+
         const forcedToolChoice = oc.systemContext.get(FORCED_TOOL_CHOICE_CONTEXT_KEY) as
           | ToolChoice<Record<string, unknown>>
           | undefined;

packages/core/src/agent/prepare-step.spec.ts

@@ -0,0 +1,122 @@
+import { describe, expect, it, vi } from "vitest";
+import { Agent } from "./agent";
+import { createMockLanguageModel, defaultMockResponse } from "./test-utils";
+
+describe("prepareStep", () => {
+  it("should accept prepareStep in AgentOptions", () => {
+    const prepareStep = vi.fn(() => ({}));
+    const model = createMockLanguageModel();
+
+    const agent = new Agent({
+      name: "test-agent",
+      instructions: "test",
+      model,
+      prepareStep,
+    });
+
+    expect(agent.prepareStep).toBe(prepareStep);
+  });
+
+  it("should default to undefined when prepareStep is not provided", () => {
+    const model = createMockLanguageModel();
+
+    const agent = new Agent({
+      name: "test-agent",
+      instructions: "test",
+      model,
+    });
+
+    expect(agent.prepareStep).toBeUndefined();
+  });
+
+  it("should pass agent-level prepareStep to generateText", async () => {
+    const prepareStep = vi.fn(() => ({}));
+    const model = createMockLanguageModel({
+      doGenerate: {
+        ...defaultMockResponse,
+        content: [{ type: "text", text: "done" }],
+      },
+    });
+
+    const agent = new Agent({
+      name: "test-agent",
+      instructions: "test",
+      model,
+      prepareStep,
+    });
+
+    await agent.generateText("hello");
+
+    // prepareStep is called by the AI SDK on each step
+    expect(prepareStep).toHaveBeenCalled();
+  });
+
+  it("should pass agent-level prepareStep to streamText", async () => {
+    const prepareStep = vi.fn(() => ({}));
+    const model = createMockLanguageModel();
+
+    const agent = new Agent({
+      name: "test-agent",
+      instructions: "test",
+      model,
+      prepareStep,
+    });
+
+    const result = await agent.streamText("hello");
+    // consume the stream to completion
+    for await (const _part of result.textStream) {
+      // drain
+    }
+
+    expect(prepareStep).toHaveBeenCalled();
+  });
+
+  it("should allow per-call prepareStep to override agent-level", async () => {
+    const agentPrepareStep = vi.fn(() => ({}));
+    const callPrepareStep = vi.fn(() => ({}));
+    const model = createMockLanguageModel({
+      doGenerate: {
+        ...defaultMockResponse,
+        content: [{ type: "text", text: "done" }],
+      },
+    });
+
+    const agent = new Agent({
+      name: "test-agent",
+      instructions: "test",
+      model,
+      prepareStep: agentPrepareStep,
+    });
+
+    await agent.generateText("hello", {
+      prepareStep: callPrepareStep,
+    });
+
+    // per-call should be used, not agent-level
+    expect(callPrepareStep).toHaveBeenCalled();
+    expect(agentPrepareStep).not.toHaveBeenCalled();
+  });
+
+  it("should allow per-call prepareStep to override agent-level in streamText", async () => {
+    const agentPrepareStep = vi.fn(() => ({}));
+    const callPrepareStep = vi.fn(() => ({}));
+    const model = createMockLanguageModel();
+
+    const agent = new Agent({
+      name: "test-agent",
+      instructions: "test",
+      model,
+      prepareStep: agentPrepareStep,
+    });
+
+    const result = await agent.streamText("hello", {
+      prepareStep: callPrepareStep,
+    });
+    for await (const _part of result.textStream) {
+      // drain
+    }
+
+    expect(callPrepareStep).toHaveBeenCalled();
+    expect(agentPrepareStep).not.toHaveBeenCalled();
+  });
+});

packages/core/src/agent/types.ts

@@ -11,7 +11,7 @@ import type {
   ProviderTextResponse,
   ProviderTextStreamResponse,
 } from "../agent/providers/base/types";
-import type { StopWhen } from "../ai-types";
+import type { PrepareStep, StopWhen } from "../ai-types";
 
 import type { LanguageModel, TextStreamPart, UIMessage } from "ai";
 import type { Memory } from "../memory";
@@ -723,6 +723,17 @@ export type AgentOptions = {
    * Per-call `stopWhen` in method options overrides this.
    */
   stopWhen?: StopWhen;
+  /**
+   * Default step preparation callback (ai-sdk `prepareStep`).
+   * Called before each step to control tool availability, tool choice, etc.
+   * Per-call `prepareStep` in method options overrides this.
+   *
+   * @example
+   * ```ts
+   * prepareStep: ({ steps }) => (steps.length > 0 ? { toolChoice: 'none' } : {}),
+   * ```
+   */
+  prepareStep?: PrepareStep;
   markdown?: boolean;
   /**
    * When true, use the active VoltAgent span as the parent if parentSpan is not provided.

packages/core/src/ai-types.ts

@@ -3,3 +3,6 @@ import type { generateText } from "ai";
 
 // StopWhen predicate type used by ai-sdk generate/stream functions
 export type StopWhen = Parameters<typeof generateText>[0]["stopWhen"];
+
+// PrepareStep callback type used by ai-sdk generate/stream functions
+export type PrepareStep = Parameters<typeof generateText>[0]["prepareStep"];

packages/core/src/index.ts

@@ -302,7 +302,7 @@ export { createAsyncIterableStream, type AsyncIterableStream } from "@voltagent/
 // Convenience re-exports from ai-sdk so apps need only @voltagent/core
 export { stepCountIs, hasToolCall } from "ai";
 export type { LanguageModel } from "ai";
-export type { StopWhen } from "./ai-types";
+export type { PrepareStep, StopWhen } from "./ai-types";
 
 export type {
   ManagedMemoryStatus,

website/docs/getting-started/migration-guide.md

@@ -610,6 +610,19 @@ console.log(out.context); // VoltAgent Map
 - This overrides VoltAgent's default `stepCountIs(maxSteps)` guard.
 - Be cautious: permissive predicates can lead to long-running or looping generations; overly strict ones may stop before tools complete.
 
+### prepareStep callback (advanced)
+
+- You can pass an ai-sdk `prepareStep` callback in `AgentOptions` or in per-call method options to control tool availability, tool choice, and other settings before each step.
+- Per-call `prepareStep` overrides the agent-level default.
+- Example: force text-only output after the first step:
+  ```ts
+  const agent = new Agent({
+    name: "my-agent",
+    model,
+    prepareStep: ({ steps }) => (steps.length > 0 ? { toolChoice: "none" } : {}),
+  });
+  ```
+
 ### Built-in server removed; use `@voltagent/server-hono`
 
 VoltAgent 1.x decouples the HTTP server from `@voltagent/core`. The built-in server is removed in favor of pluggable server providers. The recommended provider is `@voltagent/server-hono` (powered by Hono). Default port remains `3141`.