From 1c617bb5e954b2e4e09469faee48d6ab0682036b Mon Sep 17 00:00:00 2001 From: Ravish Date: Mon, 6 Apr 2026 11:20:45 -0700 Subject: [PATCH] feat(core): add prepareStep to AgentOptions for per-step tool control 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 --- .changeset/feat-agent-prepare-step.md | 9 ++ packages/core/src/agent/agent.ts | 21 ++- packages/core/src/agent/prepare-step.spec.ts | 122 ++++++++++++++++++ packages/core/src/agent/types.ts | 13 +- packages/core/src/ai-types.ts | 3 + packages/core/src/index.ts | 2 +- .../docs/getting-started/migration-guide.md | 13 ++ 7 files changed, 180 insertions(+), 3 deletions(-) create mode 100644 .changeset/feat-agent-prepare-step.md create mode 100644 packages/core/src/agent/prepare-step.spec.ts diff --git a/.changeset/feat-agent-prepare-step.md b/.changeset/feat-agent-prepare-step.md new file mode 100644 index 000000000..a9610e488 --- /dev/null +++ b/.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. diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index 1d0280c0e..1c47e6f02 100644 --- a/packages/core/src/agent/agent.ts +++ b/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>; + + /** + * 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> | 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> | undefined; diff --git a/packages/core/src/agent/prepare-step.spec.ts b/packages/core/src/agent/prepare-step.spec.ts new file mode 100644 index 000000000..e3332ba14 --- /dev/null +++ b/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(); + }); +}); diff --git a/packages/core/src/agent/types.ts b/packages/core/src/agent/types.ts index a4e6eb82c..f919cc9ec 100644 --- a/packages/core/src/agent/types.ts +++ b/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. diff --git a/packages/core/src/ai-types.ts b/packages/core/src/ai-types.ts index 12785d534..9b5e02666 100644 --- a/packages/core/src/ai-types.ts +++ b/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[0]["stopWhen"]; + +// PrepareStep callback type used by ai-sdk generate/stream functions +export type PrepareStep = Parameters[0]["prepareStep"]; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 42425f7a1..42c239ef2 100644 --- a/packages/core/src/index.ts +++ b/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, diff --git a/website/docs/getting-started/migration-guide.md b/website/docs/getting-started/migration-guide.md index 4eea7b9b4..5fe4c1ecf 100644 --- a/website/docs/getting-started/migration-guide.md +++ b/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`.