@ai-sdk/workflow 1.0.0-beta.2 → 1.0.0-beta.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ai-sdk/workflow",
-  "version": "1.0.0-beta.2",
+  "version": "1.0.0-beta.3",
   "description": "WorkflowAgent for building AI agents with AI SDK",
   "license": "Apache-2.0",
   "main": "./dist/index.js",

package/src/do-stream-step.ts

@@ -21,7 +21,6 @@ import {
   resolveSerializableTools,
   type SerializableToolDef,
 } from './serializable-schema.js';
-import type { CompatibleLanguageModel } from './types.js';
 
 export type { Experimental_LanguageModelStreamPart as ModelCallStreamPart } from 'ai';
 
@@ -87,10 +86,7 @@ export interface StreamFinish {
 
 export async function doStreamStep(
   conversationPrompt: LanguageModelV4Prompt,
-  modelInit:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>),
+  modelInit: LanguageModel,
   writable?: WritableStream<ModelCallStreamPart<ToolSet>>,
   serializedTools?: Record<string, SerializableToolDef>,
   options?: DoStreamStepOptions,
@@ -98,22 +94,10 @@ export async function doStreamStep(
   'use step';
 
   // Resolve model inside step (must happen here for serialization boundary)
-  let model: CompatibleLanguageModel;
-  if (typeof modelInit === 'string') {
-    model = gateway.languageModel(modelInit) as CompatibleLanguageModel;
-  } else if (typeof modelInit === 'function') {
-    model = await modelInit();
-  } else if (
-    typeof modelInit === 'object' &&
-    modelInit !== null &&
-    'modelId' in modelInit
-  ) {
-    model = modelInit;
-  } else {
-    throw new Error(
-      'Invalid "model initialization" argument. Must be a string, a LanguageModel instance, or a function that returns a LanguageModel instance.',
-    );
-  }
+  const model: LanguageModel =
+    typeof modelInit === 'string'
+      ? gateway.languageModel(modelInit)
+      : modelInit;
 
   // Reconstruct tools from serializable definitions with Ajv validation.
   // Tools are serialized before crossing the step boundary because zod schemas
@@ -126,7 +110,7 @@ export async function doStreamStep(
   // model.doStream(), retry logic, and stream part transformation
   // (tool call parsing, finish reason mapping, file wrapping).
   const { stream: modelStream } = await streamModelCall({
-    model: model as LanguageModel,
+    model,
     // streamModelCall expects Prompt (ModelMessage[]) but we pass the
     // pre-converted LanguageModelV4Prompt. standardizePrompt inside
     // streamModelCall handles both formats.

package/src/index.ts

@@ -17,9 +17,10 @@ export {
   type PrepareStepInfo,
   type PrepareStepResult,
   type ProviderOptions,
-  type StreamTextOnAbortCallback,
-  type StreamTextOnErrorCallback,
-  type StreamTextOnFinishCallback,
+  type WorkflowAgentOnAbortCallback,
+  type WorkflowAgentOnErrorCallback,
+  type WorkflowAgentOnFinishCallback,
+  type WorkflowAgentOnStepFinishCallback,
   type StreamTextTransform,
   type TelemetrySettings,
   type ToolCallRepairFunction,

package/src/providers/mock.ts

@@ -6,19 +6,55 @@ export type MockResponseDescriptor =
 
 /**
  * Mock model that returns a fixed text response.
- * Same 'use step' pattern as real providers (anthropic, openai, etc.).
- * Only captures `text` (string) — fully serializable across step boundary.
  */
 export function mockTextModel(text: string) {
-  return async () => {
-    'use step';
-    // Bind closure var at step body level so SWC plugin detects it
-    const _text = text;
-    return mockProvider({
-      doStream: async () => ({
-        stream: new ReadableStream({
-          start(c) {
-            for (const v of [
+  return mockProvider({
+    doStream: async () => ({
+      stream: new ReadableStream({
+        start(c) {
+          for (const v of [
+            { type: 'stream-start', warnings: [] },
+            {
+              type: 'response-metadata',
+              id: 'r',
+              modelId: 'mock',
+              timestamp: new Date(),
+            },
+            { type: 'text-start', id: '1' },
+            { type: 'text-delta', id: '1', delta: text },
+            { type: 'text-end', id: '1' },
+            {
+              type: 'finish',
+              finishReason: { unified: 'stop', raw: 'stop' },
+              usage: {
+                inputTokens: { total: 5, noCache: 5 },
+                outputTokens: { total: 10, text: 10 },
+              },
+            },
+          ] as any[])
+            c.enqueue(v);
+          c.close();
+        },
+      }),
+    }),
+  });
+}
+
+/**
+ * Mock model that plays through a sequence of responses.
+ * Determines which response to return by counting assistant messages in the prompt.
+ */
+export function mockSequenceModel(responses: MockResponseDescriptor[]) {
+  return mockProvider({
+    doStream: async (options: any) => {
+      const idx = Math.min(
+        options.prompt.filter((m: any) => m.role === 'assistant').length,
+        responses.length - 1,
+      );
+      const r = responses[idx];
+      const parts =
+        r.type === 'text'
+          ? [
               { type: 'stream-start', warnings: [] },
               {
                 type: 'response-metadata',
@@ -27,7 +63,7 @@ export function mockTextModel(text: string) {
                 timestamp: new Date(),
               },
               { type: 'text-start', id: '1' },
-              { type: 'text-delta', id: '1', delta: _text },
+              { type: 'text-delta', id: '1', delta: r.text },
               { type: 'text-end', id: '1' },
               {
                 type: 'finish',
@@ -37,87 +73,38 @@ export function mockTextModel(text: string) {
                   outputTokens: { total: 10, text: 10 },
                 },
               },
-            ] as any[])
-              c.enqueue(v);
+            ]
+          : [
+              { type: 'stream-start', warnings: [] },
+              {
+                type: 'response-metadata',
+                id: 'r',
+                modelId: 'mock',
+                timestamp: new Date(),
+              },
+              {
+                type: 'tool-call',
+                toolCallId: `call-${idx + 1}`,
+                toolName: r.toolName,
+                input: r.input,
+              },
+              {
+                type: 'finish',
+                finishReason: { unified: 'tool-calls', raw: undefined },
+                usage: {
+                  inputTokens: { total: 5, noCache: 5 },
+                  outputTokens: { total: 10, text: 10 },
+                },
+              },
+            ];
+      return {
+        stream: new ReadableStream({
+          start(c) {
+            for (const p of parts as any[]) c.enqueue(p);
             c.close();
           },
         }),
-      }),
-    });
-  };
-}
-
-/**
- * Mock model that plays through a sequence of responses.
- * Determines which response to return by counting assistant messages in the prompt.
- * Only captures `responses` (array of plain objects) — fully serializable.
- */
-export function mockSequenceModel(responses: MockResponseDescriptor[]) {
-  return async () => {
-    'use step';
-    // Bind closure var at step body level so SWC plugin detects it
-    const _responses = responses;
-    return mockProvider({
-      doStream: async (options: any) => {
-        const idx = Math.min(
-          options.prompt.filter((m: any) => m.role === 'assistant').length,
-          _responses.length - 1,
-        );
-        const r = _responses[idx];
-        const parts =
-          r.type === 'text'
-            ? [
-                { type: 'stream-start', warnings: [] },
-                {
-                  type: 'response-metadata',
-                  id: 'r',
-                  modelId: 'mock',
-                  timestamp: new Date(),
-                },
-                { type: 'text-start', id: '1' },
-                { type: 'text-delta', id: '1', delta: r.text },
-                { type: 'text-end', id: '1' },
-                {
-                  type: 'finish',
-                  finishReason: { unified: 'stop', raw: 'stop' },
-                  usage: {
-                    inputTokens: { total: 5, noCache: 5 },
-                    outputTokens: { total: 10, text: 10 },
-                  },
-                },
-              ]
-            : [
-                { type: 'stream-start', warnings: [] },
-                {
-                  type: 'response-metadata',
-                  id: 'r',
-                  modelId: 'mock',
-                  timestamp: new Date(),
-                },
-                {
-                  type: 'tool-call',
-                  toolCallId: `call-${idx + 1}`,
-                  toolName: r.toolName,
-                  input: r.input,
-                },
-                {
-                  type: 'finish',
-                  finishReason: { unified: 'tool-calls', raw: undefined },
-                  usage: {
-                    inputTokens: { total: 5, noCache: 5 },
-                    outputTokens: { total: 10, text: 10 },
-                  },
-                },
-              ];
-        return {
-          stream: new ReadableStream({
-            start(c) {
-              for (const p of parts as any[]) c.enqueue(p);
-              c.close();
-            },
-          }),
-        };
-      },
-    });
-  };
+      };
+    },
+  });
 }

package/src/stream-text-iterator.ts

@@ -5,9 +5,9 @@ import type {
 } from '@ai-sdk/provider';
 import type {
   Experimental_LanguageModelStreamPart as ModelCallStreamPart,
+  LanguageModel,
   ModelMessage,
   StepResult,
-  StreamTextOnStepFinishCallback,
   ToolCallRepairFunction,
   ToolChoice,
   ToolSet,
@@ -22,11 +22,11 @@ import { serializeToolSet } from './serializable-schema.js';
 import type {
   GenerationSettings,
   PrepareStepCallback,
-  StreamTextOnErrorCallback,
+  WorkflowAgentOnErrorCallback,
+  WorkflowAgentOnStepFinishCallback,
   TelemetrySettings,
   WorkflowAgentOnStepStartCallback,
 } from './workflow-agent.js';
-import type { CompatibleLanguageModel } from './types.js';
 
 // Re-export for consumers
 export type { ProviderExecutedToolResult } from './do-stream-step.js';
@@ -71,15 +71,12 @@ export async function* streamTextIterator({
   prompt: LanguageModelV4Prompt;
   tools: ToolSet;
   writable?: WritableStream<ModelCallStreamPart<ToolSet>>;
-  model:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  model: LanguageModel;
   stopConditions?: ModelStopCondition[] | ModelStopCondition;
   maxSteps?: number;
-  onStepFinish?: StreamTextOnStepFinishCallback<any, any>;
+  onStepFinish?: WorkflowAgentOnStepFinishCallback<any>;
   onStepStart?: WorkflowAgentOnStepStartCallback;
-  onError?: StreamTextOnErrorCallback;
+  onError?: WorkflowAgentOnErrorCallback;
   prepareStep?: PrepareStepCallback<any>;
   generationSettings?: GenerationSettings;
   toolChoice?: ToolChoice<ToolSet>;
@@ -94,10 +91,7 @@ export async function* streamTextIterator({
   LanguageModelV4ToolResultPart[]
 > {
   let conversationPrompt = [...prompt]; // Create a mutable copy
-  let currentModel:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>) = model;
+  let currentModel: LanguageModel = model;
   let currentGenerationSettings = generationSettings ?? {};
   let currentToolChoice = toolChoice;
   let currentContext = experimental_context;

package/src/workflow-agent.ts

@@ -21,6 +21,7 @@ import {
   type ToolChoice,
   type ToolSet,
   type UIMessage,
+  LanguageModel,
 } from 'ai';
 import {
   convertToLanguageModelPrompt,
@@ -34,6 +35,15 @@ import type { CompatibleLanguageModel } from './types.js';
 // Re-export for consumers
 export type { CompatibleLanguageModel } from './types.js';
 
+/**
+ * Callback function to be called after each step completes.
+ * Alias for the AI SDK's StreamTextOnStepFinishCallback, using
+ * WorkflowAgent-consistent naming.
+ */
+export type WorkflowAgentOnStepFinishCallback<
+  TTools extends ToolSet = ToolSet,
+> = StreamTextOnStepFinishCallback<TTools, any>;
+
 /**
  * Infer the type of the tools of a workflow agent.
  */
@@ -248,10 +258,7 @@ export interface PrepareStepInfo<TTools extends ToolSet = ToolSet> {
    * The current model configuration (string or function).
    * The function should return a LanguageModelV4 instance.
    */
-  model:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  model: LanguageModel;
 
   /**
    * The current step number (0-indexed).
@@ -282,12 +289,8 @@ export interface PrepareStepInfo<TTools extends ToolSet = ToolSet> {
 export interface PrepareStepResult extends Partial<GenerationSettings> {
   /**
    * Override the model for this step.
-   * The function should return a LanguageModelV4 instance.
    */
-  model?:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  model?: LanguageModel;
 
   /**
    * Override the system message for this step.
@@ -332,10 +335,7 @@ export type PrepareStepCallback<TTools extends ToolSet = ToolSet> = (
 export interface PrepareCallOptions<
   TTools extends ToolSet = ToolSet,
 > extends Partial<GenerationSettings> {
-  model:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  model: LanguageModel;
   tools: TTools;
   instructions?: string | SystemModelMessage | Array<SystemModelMessage>;
   toolChoice?: ToolChoice<TTools>;
@@ -371,12 +371,9 @@ export interface WorkflowAgentOptions<
    * The model provider to use for the agent.
    *
    * This should be a string compatible with the Vercel AI Gateway (e.g., 'anthropic/claude-opus'),
-   * or a step function that returns a LanguageModelV4 instance.
+   * or a LanguageModelV4 instance from a provider.
    */
-  model:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  model: LanguageModel;
 
   /**
    * A set of tools available to the agent.
@@ -428,12 +425,12 @@ export interface WorkflowAgentOptions<
   /**
    * Callback function to be called after each step completes.
    */
-  onStepFinish?: StreamTextOnStepFinishCallback<ToolSet, any>;
+  onStepFinish?: WorkflowAgentOnStepFinishCallback<ToolSet>;
 
   /**
    * Callback that is called when the LLM response and all request tool executions are finished.
    */
-  onFinish?: StreamTextOnFinishCallback<ToolSet>;
+  onFinish?: WorkflowAgentOnFinishCallback<ToolSet>;
 
   /**
    * Callback called when the agent starts streaming, before any LLM calls.
@@ -466,7 +463,7 @@ export interface WorkflowAgentOptions<
 /**
  * Callback that is called when the LLM response and all request tool executions are finished.
  */
-export type StreamTextOnFinishCallback<
+export type WorkflowAgentOnFinishCallback<
   TTools extends ToolSet = ToolSet,
   OUTPUT = never,
 > = (event: {
@@ -510,14 +507,14 @@ export type StreamTextOnFinishCallback<
 /**
  * Callback that is invoked when an error occurs during streaming.
  */
-export type StreamTextOnErrorCallback = (event: {
+export type WorkflowAgentOnErrorCallback = (event: {
   error: unknown;
 }) => PromiseLike<void> | void;
 
 /**
  * Callback that is set using the `onAbort` option.
  */
-export type StreamTextOnAbortCallback<TTools extends ToolSet = ToolSet> =
+export type WorkflowAgentOnAbortCallback<TTools extends ToolSet = ToolSet> =
   (event: {
     /**
      * Details for all previously finished steps.
@@ -530,10 +527,7 @@ export type StreamTextOnAbortCallback<TTools extends ToolSet = ToolSet> =
  */
 export type WorkflowAgentOnStartCallback = (event: {
   /** The model being used */
-  readonly model:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  readonly model: LanguageModel;
   /** The messages being sent */
   readonly messages: ModelMessage[];
 }) => PromiseLike<void> | void;
@@ -545,10 +539,7 @@ export type WorkflowAgentOnStepStartCallback = (event: {
   /** The current step number (0-based) */
   readonly stepNumber: number;
   /** The model being used for this step */
-  readonly model:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  readonly model: LanguageModel;
   /** The messages being sent for this step */
   readonly messages: ModelMessage[];
 }) => PromiseLike<void> | void;
@@ -710,24 +701,24 @@ export interface WorkflowAgentStreamOptions<
   /**
    * Callback function to be called after each step completes.
    */
-  onStepFinish?: StreamTextOnStepFinishCallback<TTools, any>;
+  onStepFinish?: WorkflowAgentOnStepFinishCallback<TTools>;
 
   /**
    * Callback that is invoked when an error occurs during streaming.
    * You can use it to log errors.
    */
-  onError?: StreamTextOnErrorCallback;
+  onError?: WorkflowAgentOnErrorCallback;
 
   /**
    * Callback that is called when the LLM response and all request tool executions
    * (for tools that have an `execute` function) are finished.
    */
-  onFinish?: StreamTextOnFinishCallback<TTools, OUTPUT>;
+  onFinish?: WorkflowAgentOnFinishCallback<TTools, OUTPUT>;
 
   /**
    * Callback that is called when the operation is aborted.
    */
-  onAbort?: StreamTextOnAbortCallback<TTools>;
+  onAbort?: WorkflowAgentOnAbortCallback<TTools>;
 
   /**
    * Callback called when the agent starts streaming, before any LLM calls.
@@ -898,10 +889,7 @@ export interface WorkflowAgentStreamResult<
  * ```
  */
 export class WorkflowAgent<TBaseTools extends ToolSet = ToolSet> {
-  private model:
-    | string
-    | CompatibleLanguageModel
-    | (() => Promise<CompatibleLanguageModel>);
+  private model: LanguageModel;
   /**
    * The tool set configured for this agent.
    */
@@ -915,11 +903,8 @@ export class WorkflowAgent<TBaseTools extends ToolSet = ToolSet> {
   private telemetry?: TelemetrySettings;
   private experimentalContext: unknown;
   private prepareStep?: PrepareStepCallback<TBaseTools>;
-  private constructorOnStepFinish?: StreamTextOnStepFinishCallback<
-    ToolSet,
-    any
-  >;
-  private constructorOnFinish?: StreamTextOnFinishCallback<ToolSet>;
+  private constructorOnStepFinish?: WorkflowAgentOnStepFinishCallback<ToolSet>;
+  private constructorOnFinish?: WorkflowAgentOnFinishCallback<ToolSet>;
   private constructorOnStart?: WorkflowAgentOnStartCallback;
   private constructorOnStepStart?: WorkflowAgentOnStepStartCallback;
   private constructorOnToolCallStart?: WorkflowAgentOnToolCallStartCallback;
@@ -972,10 +957,7 @@ export class WorkflowAgent<TBaseTools extends ToolSet = ToolSet> {
     options: WorkflowAgentStreamOptions<TTools, OUTPUT, PARTIAL_OUTPUT>,
   ): Promise<WorkflowAgentStreamResult<TTools, OUTPUT>> {
     // Call prepareCall to transform parameters before the agent loop
-    let effectiveModel:
-      | string
-      | CompatibleLanguageModel
-      | (() => Promise<CompatibleLanguageModel>) = this.model;
+    let effectiveModel: LanguageModel = this.model;
     let effectiveInstructions = options.system ?? this.instructions;
     let effectiveMessages = options.messages;
     let effectiveGenerationSettings = { ...this.generationSettings };
@@ -1210,13 +1192,13 @@ export class WorkflowAgent<TBaseTools extends ToolSet = ToolSet> {
     // Merge constructor + stream callbacks (constructor first, then stream)
     const mergedOnStepFinish = mergeCallbacks(
       this.constructorOnStepFinish as
-        | StreamTextOnStepFinishCallback<TTools, any>
+        | WorkflowAgentOnStepFinishCallback<TTools>
         | undefined,
       options.onStepFinish,
     );
     const mergedOnFinish = mergeCallbacks(
       this.constructorOnFinish as
-        | StreamTextOnFinishCallback<TTools, OUTPUT>
+        | WorkflowAgentOnFinishCallback<TTools, OUTPUT>
         | undefined,
       options.onFinish,
     );