@ai-sdk/workflow 1.0.0-beta.8 → 1.0.0-beta.94

package/src/stream-text-iterator.ts

@@ -3,15 +3,18 @@ import type {
   LanguageModelV4Prompt,
   LanguageModelV4ToolResultPart,
 } from '@ai-sdk/provider';
-import type {
-  Experimental_LanguageModelStreamPart as ModelCallStreamPart,
-  LanguageModel,
-  ModelMessage,
-  StepResult,
-  ToolCallRepairFunction,
-  ToolChoice,
-  ToolSet,
+import type { Context } from '@ai-sdk/provider-utils';
+import {
+  experimental_filterActiveTools as filterActiveTools,
+  type Experimental_LanguageModelStreamPart as ModelCallStreamPart,
+  type LanguageModel,
+  type ModelMessage,
+  type StepResult,
+  type ToolCallRepairFunction,
+  type ToolChoice,
+  type ToolSet,
 } from 'ai';
+import { createRestrictedTelemetryDispatcher } from 'ai/internal';
 import {
   doStreamStep,
   type ModelStopCondition,
@@ -23,8 +26,9 @@ import type {
   GenerationSettings,
   PrepareStepCallback,
   WorkflowAgentOnErrorCallback,
+  WorkflowAgentOnStepEndCallback,
   WorkflowAgentOnStepFinishCallback,
-  TelemetrySettings,
+  TelemetryOptions,
   WorkflowAgentOnStepStartCallback,
 } from './workflow-agent.js';
 
@@ -42,8 +46,10 @@ export interface StreamTextIteratorYieldValue {
   messages: LanguageModelV4Prompt;
   /** The step result from the current step */
   step?: StepResult<ToolSet, any>;
-  /** The current experimental context */
-  context?: unknown;
+  /** The current runtime context shared across the agent loop */
+  runtimeContext?: Context;
+  /** The current per-tool context, keyed by tool name */
+  toolsContext?: Record<string, Context | undefined>;
   /** Provider-executed tool results (keyed by tool call ID) */
   providerExecutedToolResults?: Map<string, ProviderExecutedToolResult>;
 }
@@ -55,15 +61,16 @@ export async function* streamTextIterator({
   writable,
   model,
   stopConditions,
-  maxSteps,
+  onStepEnd,
   onStepFinish,
   onStepStart,
   onError,
   prepareStep,
   generationSettings,
   toolChoice,
-  experimental_context,
-  experimental_telemetry,
+  runtimeContext,
+  toolsContext,
+  telemetry,
   includeRawChunks = false,
   repairToolCall,
   responseFormat,
@@ -73,15 +80,17 @@ export async function* streamTextIterator({
   writable?: WritableStream<ModelCallStreamPart<ToolSet>>;
   model: LanguageModel;
   stopConditions?: ModelStopCondition[] | ModelStopCondition;
-  maxSteps?: number;
+  onStepEnd?: WorkflowAgentOnStepEndCallback<any>;
+  /** @deprecated Use `onStepEnd` instead. */
   onStepFinish?: WorkflowAgentOnStepFinishCallback<any>;
   onStepStart?: WorkflowAgentOnStepStartCallback;
   onError?: WorkflowAgentOnErrorCallback;
   prepareStep?: PrepareStepCallback<any>;
   generationSettings?: GenerationSettings;
   toolChoice?: ToolChoice<ToolSet>;
-  experimental_context?: unknown;
-  experimental_telemetry?: TelemetrySettings;
+  runtimeContext?: Context;
+  toolsContext?: Record<string, Context | undefined>;
+  telemetry?: TelemetryOptions<Context, ToolSet>;
   includeRawChunks?: boolean;
   repairToolCall?: ToolCallRepairFunction<ToolSet>;
   responseFormat?: LanguageModelV4CallOptions['responseFormat'];
@@ -94,7 +103,9 @@ export async function* streamTextIterator({
   let currentModel: LanguageModel = model;
   let currentGenerationSettings = generationSettings ?? {};
   let currentToolChoice = toolChoice;
-  let currentContext = experimental_context;
+  let currentRuntimeContext: Context = runtimeContext ?? {};
+  let currentToolsContext: Record<string, Context | undefined> =
+    toolsContext ?? {};
   let currentActiveTools: string[] | undefined;
 
   const steps: StepResult<any, any>[] = [];
@@ -104,16 +115,21 @@ export async function* streamTextIterator({
   let lastStep: StepResult<any, any> | undefined;
   let lastStepWasToolCalls = false;
 
-  // Default maxSteps to Infinity to preserve backwards compatibility
-  // (agent loops until completion unless explicitly limited)
-  const effectiveMaxSteps = maxSteps ?? Infinity;
+  // TODO(#12164): replace this AI-core telemetry bridge with a
+  // WorkflowAgent-specific typed dispatcher. `streamTextIterator` widens
+  // tools/runtime context and emits Workflow-shaped events that are only
+  // approximately compatible with generateText telemetry event types.
+  const telemetryDispatcher = createRestrictedTelemetryDispatcher<
+    any,
+    any,
+    any
+  >({
+    telemetry: telemetry as any,
+    includeRuntimeContext: telemetry?.includeRuntimeContext,
+    includeToolsContext: telemetry?.includeToolsContext,
+  }) as any;
 
   while (!done) {
-    // Check if we've exceeded the maximum number of steps
-    if (stepNumber >= effectiveMaxSteps) {
-      break;
-    }
-
     // Check for abort signal
     if (currentGenerationSettings.abortSignal?.aborted) {
       break;
@@ -126,19 +142,20 @@ export async function* streamTextIterator({
         stepNumber,
         steps,
         messages: conversationPrompt,
-        experimental_context: currentContext,
+        runtimeContext: currentRuntimeContext,
+        toolsContext: currentToolsContext as never,
       });
 
       // Apply any overrides from prepareStep
-      if (prepareResult.model !== undefined) {
+      if (prepareResult?.model !== undefined) {
         currentModel = prepareResult.model;
       }
       // Apply messages override BEFORE system so the system message
       // isn't lost when messages replaces the prompt.
-      if (prepareResult.messages !== undefined) {
+      if (prepareResult?.messages !== undefined) {
         conversationPrompt = [...prepareResult.messages];
       }
-      if (prepareResult.system !== undefined) {
+      if (prepareResult?.system !== undefined) {
         // Update or prepend system message in the conversation prompt.
         // Applied AFTER messages override so the system message isn't
         // lost when messages replaces the prompt.
@@ -159,80 +176,86 @@ export async function* streamTextIterator({
           });
         }
       }
-      if (prepareResult.experimental_context !== undefined) {
-        currentContext = prepareResult.experimental_context;
+      if (prepareResult?.runtimeContext !== undefined) {
+        currentRuntimeContext = prepareResult.runtimeContext;
       }
-      if (prepareResult.activeTools !== undefined) {
+      if (prepareResult?.toolsContext !== undefined) {
+        currentToolsContext = prepareResult.toolsContext as Record<
+          string,
+          Context | undefined
+        >;
+      }
+      if (prepareResult?.activeTools !== undefined) {
         currentActiveTools = prepareResult.activeTools;
       }
       // Apply generation settings overrides
-      if (prepareResult.maxOutputTokens !== undefined) {
+      if (prepareResult?.maxOutputTokens !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           maxOutputTokens: prepareResult.maxOutputTokens,
         };
       }
-      if (prepareResult.temperature !== undefined) {
+      if (prepareResult?.temperature !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           temperature: prepareResult.temperature,
         };
       }
-      if (prepareResult.topP !== undefined) {
+      if (prepareResult?.topP !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           topP: prepareResult.topP,
         };
       }
-      if (prepareResult.topK !== undefined) {
+      if (prepareResult?.topK !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           topK: prepareResult.topK,
         };
       }
-      if (prepareResult.presencePenalty !== undefined) {
+      if (prepareResult?.presencePenalty !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           presencePenalty: prepareResult.presencePenalty,
         };
       }
-      if (prepareResult.frequencyPenalty !== undefined) {
+      if (prepareResult?.frequencyPenalty !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           frequencyPenalty: prepareResult.frequencyPenalty,
         };
       }
-      if (prepareResult.stopSequences !== undefined) {
+      if (prepareResult?.stopSequences !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           stopSequences: prepareResult.stopSequences,
         };
       }
-      if (prepareResult.seed !== undefined) {
+      if (prepareResult?.seed !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           seed: prepareResult.seed,
         };
       }
-      if (prepareResult.maxRetries !== undefined) {
+      if (prepareResult?.maxRetries !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           maxRetries: prepareResult.maxRetries,
         };
       }
-      if (prepareResult.headers !== undefined) {
+      if (prepareResult?.headers !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           headers: prepareResult.headers,
         };
       }
-      if (prepareResult.providerOptions !== undefined) {
+      if (prepareResult?.providerOptions !== undefined) {
         currentGenerationSettings = {
           ...currentGenerationSettings,
           providerOptions: prepareResult.providerOptions,
         };
       }
-      if (prepareResult.toolChoice !== undefined) {
+      if (prepareResult?.toolChoice !== undefined) {
         currentToolChoice = prepareResult.toolChoice;
       }
     }
@@ -243,20 +266,66 @@ export async function* streamTextIterator({
         model: currentModel,
         messages: conversationPrompt as unknown as ModelMessage[],
         steps: [...steps],
+        runtimeContext: currentRuntimeContext,
+        toolsContext: currentToolsContext as never,
       });
     }
 
+    const stepStartModelInfo = getModelInfo(currentModel);
+    await telemetryDispatcher.onStepStart?.({
+      callId: 'workflow-agent',
+      provider: stepStartModelInfo.provider,
+      modelId: stepStartModelInfo.modelId,
+      stepNumber,
+      system: undefined,
+      messages: conversationPrompt as unknown as ModelMessage[],
+      tools,
+      toolChoice: currentToolChoice,
+      activeTools: currentActiveTools as never,
+      steps: steps.map(normalizeStepForTelemetry),
+      providerOptions: currentGenerationSettings.providerOptions,
+      output: undefined,
+      runtimeContext: currentRuntimeContext,
+      toolsContext: currentToolsContext as never,
+    });
+
     try {
       // Filter tools if activeTools is specified
       const effectiveTools =
         currentActiveTools && currentActiveTools.length > 0
-          ? filterToolSet(tools, currentActiveTools)
+          ? (filterActiveTools({
+              tools,
+              activeTools: currentActiveTools,
+            }) ?? tools)
           : tools;
 
       // Serialize tools before crossing the step boundary — zod schemas
       // contain functions that can't be serialized by the workflow runtime.
       // Tools are reconstructed with Ajv validation inside doStreamStep.
       const serializedTools = serializeToolSet(effectiveTools);
+      const modelCallInfo = getModelInfo(currentModel);
+
+      await telemetryDispatcher.onLanguageModelCallStart?.({
+        callId: 'workflow-agent',
+        provider: modelCallInfo.provider,
+        modelId: modelCallInfo.modelId,
+        system: undefined,
+        messages: conversationPrompt as unknown as ModelMessage[],
+        tools:
+          serializedTools == null
+            ? undefined
+            : Object.values(serializedTools).map(tool => ({ ...tool })),
+        maxOutputTokens: currentGenerationSettings.maxOutputTokens,
+        temperature: currentGenerationSettings.temperature,
+        topP: currentGenerationSettings.topP,
+        topK: currentGenerationSettings.topK,
+        presencePenalty: currentGenerationSettings.presencePenalty,
+        frequencyPenalty: currentGenerationSettings.frequencyPenalty,
+        stopSequences: currentGenerationSettings.stopSequences,
+        seed: currentGenerationSettings.seed,
+        providerOptions: currentGenerationSettings.providerOptions,
+        headers: currentGenerationSettings.headers,
+      } as never);
 
       const { toolCalls, finish, step, providerExecutedToolResults } =
         await doStreamStep(
@@ -268,12 +337,24 @@ export async function* streamTextIterator({
             ...currentGenerationSettings,
             toolChoice: currentToolChoice,
             includeRawChunks,
-            experimental_telemetry,
             repairToolCall,
             responseFormat,
+            runtimeContext: currentRuntimeContext,
+            toolsContext: currentToolsContext,
+            stepNumber,
           },
         );
 
+      await telemetryDispatcher.onLanguageModelCallEnd?.({
+        callId: step.callId,
+        provider: step.model?.provider ?? 'unknown',
+        modelId: step.model?.modelId ?? 'unknown',
+        finishReason: step.finishReason,
+        usage: step.usage,
+        content: step.content,
+        responseId: step.response.id,
+      });
+
       _isFirstIteration = false;
       stepNumber++;
       steps.push(step);
@@ -315,7 +396,8 @@ export async function* streamTextIterator({
           toolCalls,
           messages: conversationPrompt,
           step,
-          context: currentContext,
+          runtimeContext: currentRuntimeContext,
+          toolsContext: currentToolsContext,
           providerExecutedToolResults,
         };
 
@@ -370,9 +452,11 @@ export async function* streamTextIterator({
         );
       }
 
-      if (onStepFinish) {
-        await onStepFinish(step);
+      const resolvedOnStepEnd = onStepEnd ?? onStepFinish;
+      if (resolvedOnStepEnd) {
+        await resolvedOnStepEnd(step);
       }
+      await telemetryDispatcher.onStepEnd?.(normalizeStepForTelemetry(step));
     } catch (error) {
       if (onError) {
         await onError({ error });
@@ -387,24 +471,28 @@ export async function* streamTextIterator({
       toolCalls: [],
       messages: conversationPrompt,
       step: lastStep,
-      context: currentContext,
+      runtimeContext: currentRuntimeContext,
+      toolsContext: currentToolsContext,
     };
   }
 
   return conversationPrompt;
 }
 
-/**
- * Filter a tool set to only include the specified active tools.
- */
-function filterToolSet(tools: ToolSet, activeTools: string[]): ToolSet {
-  const filtered: ToolSet = {};
-  for (const toolName of activeTools) {
-    if (toolName in tools) {
-      filtered[toolName] = tools[toolName];
-    }
-  }
-  return filtered;
+function getModelInfo(model: LanguageModel): {
+  provider: string;
+  modelId: string;
+} {
+  return typeof model === 'string'
+    ? { provider: model.split('/')[0] ?? 'gateway', modelId: model }
+    : { provider: model.provider, modelId: model.modelId };
+}
+
+function normalizeStepForTelemetry(step: StepResult<any, any>) {
+  return {
+    ...step,
+    model: step.model ?? { provider: 'unknown', modelId: 'unknown' },
+  };
 }
 
 /**

package/src/test/agent-e2e-workflows.ts

@@ -5,7 +5,7 @@ import { tool } from 'ai';
 import { WorkflowAgent } from '../workflow-agent.js';
 import { mockTextModel, mockSequenceModel } from '../providers/mock.js';
 import { FatalError, getWritable } from 'workflow';
-import z from 'zod';
+import { z } from 'zod';
 
 // ============================================================================
 // Tool step functions
@@ -334,10 +334,10 @@ export async function agentOnStepStartE2e() {
 }
 
 // ============================================================================
-// GAP tests — experimental_onToolCallStart
+// GAP tests — onToolExecutionStart
 // ============================================================================
 
-export async function agentOnToolCallStartE2e() {
+export async function agentonToolExecutionStartE2e() {
   'use workflow';
   const calls: string[] = [];
   const agent = new WorkflowAgent({
@@ -356,14 +356,14 @@ export async function agentOnToolCallStartE2e() {
         execute: echoStep,
       },
     },
-    experimental_onToolCallStart: async () => {
+    onToolExecutionStart: async () => {
       calls.push('constructor');
     },
   } as any);
   await agent.stream({
     messages: [{ role: 'user', content: 'test' }],
     writable: getWritable(),
-    experimental_onToolCallStart: async () => {
+    onToolExecutionStart: async () => {
       calls.push('method');
     },
   } as any);
@@ -371,10 +371,10 @@ export async function agentOnToolCallStartE2e() {
 }
 
 // ============================================================================
-// GAP tests — experimental_onToolCallFinish
+// GAP tests — onToolExecutionEnd
 // ============================================================================
 
-export async function agentOnToolCallFinishE2e() {
+export async function agentonToolExecutionEndE2e() {
   'use workflow';
   const calls: string[] = [];
   let capturedEvent: any = null;
@@ -394,14 +394,14 @@ export async function agentOnToolCallFinishE2e() {
         execute: addNumbers,
       },
     },
-    experimental_onToolCallFinish: async () => {
+    onToolExecutionEnd: async () => {
       calls.push('constructor');
     },
   } as any);
   await agent.stream({
     messages: [{ role: 'user', content: 'test' }],
     writable: getWritable(),
-    experimental_onToolCallFinish: async (event: any) => {
+    onToolExecutionEnd: async (event: any) => {
       calls.push('method');
       capturedEvent = {
         toolName: event?.toolCall?.toolName,
@@ -505,3 +505,83 @@ export async function agentToolInputSchemaE2e(a: number, b: number) {
     lastStepText: result.steps[result.steps.length - 1]?.text,
   };
 }
+
+// ============================================================================
+// runtimeContext + toolsContext (end-to-end)
+// ============================================================================
+
+/**
+ * Demonstrates the full context flow:
+ *
+ * - `runtimeContext` holds shared agent state (`tenantId`, `requestId`).
+ *   `prepareStep` reads it and tags it with the current step number;
+ *   `onFinish` receives the final value.
+ * - `toolsContext` holds per-tool, schema-validated context. The
+ *   `lookupCustomer` tool declares `contextSchema`, so its entry is
+ *   validated and the tool's `execute` only sees its own context.
+ */
+export async function agentRuntimeAndToolsContextE2e() {
+  'use workflow';
+
+  let onFinishRuntimeContext: Record<string, unknown> | undefined;
+  let onFinishToolsContext: Record<string, unknown> | undefined;
+  let toolReceivedContext: unknown;
+
+  const agent = new WorkflowAgent({
+    model: mockSequenceModel([
+      {
+        type: 'tool-call',
+        toolName: 'lookupCustomer',
+        input: JSON.stringify({ customerId: 'cust_123' }),
+      },
+      { type: 'text', text: 'Customer cust_123 is eligible.' },
+    ]),
+    tools: {
+      lookupCustomer: tool({
+        description: 'Look up customer account details.',
+        inputSchema: z.object({ customerId: z.string() }),
+        contextSchema: z.object({
+          apiKey: z.string(),
+          region: z.enum(['us', 'eu']),
+        }),
+        execute: async (input, { context }) => {
+          toolReceivedContext = context;
+          return { customerId: input.customerId, eligible: true };
+        },
+      }),
+    },
+    instructions: 'You look up customers.',
+    runtimeContext: {
+      tenantId: 'tenant_123',
+      requestId: 'req_abc',
+    },
+    toolsContext: {
+      lookupCustomer: {
+        apiKey: 'sk-test-key',
+        region: 'us',
+      },
+    },
+    prepareStep: ({ stepNumber, runtimeContext }) => ({
+      runtimeContext: { ...runtimeContext, lastStep: stepNumber },
+    }),
+    onFinish: ({ runtimeContext, toolsContext }) => {
+      onFinishRuntimeContext = runtimeContext;
+      onFinishToolsContext = toolsContext;
+    },
+  });
+
+  const result = await agent.stream({
+    messages: [
+      { role: 'user', content: 'Is customer cust_123 eligible for support?' },
+    ],
+    writable: getWritable(),
+  });
+
+  return {
+    stepCount: result.steps.length,
+    lastStepText: result.steps[result.steps.length - 1]?.text,
+    toolReceivedContext,
+    onFinishRuntimeContext,
+    onFinishToolsContext,
+  };
+}

package/src/to-ui-message-chunk.ts

@@ -1,5 +1,8 @@
-import { type ToolSet, type UIMessageChunk } from 'ai';
-import type { Experimental_LanguageModelStreamPart as ModelCallStreamPart } from 'ai';
+import type {
+  Experimental_LanguageModelStreamPart as ModelCallStreamPart,
+  ToolSet,
+  UIMessageChunk,
+} from 'ai';
 
 /**
  * Convert a single ModelCallStreamPart to a UIMessageChunk.
@@ -188,20 +191,20 @@ export function toUIMessageChunk(
       // standard ModelCallStreamPart types but are written by the
       // WorkflowAgent between tool execution and the next model step
       // to ensure proper message splitting in convertToModelMessages.
-      const p = part as any;
-      if (p.type === 'tool-approval-request') {
+      const passthroughPart = part as any;
+      if (passthroughPart.type === 'tool-approval-request') {
         return {
           type: 'tool-approval-request',
-          approvalId: p.approvalId,
-          toolCallId: p.toolCallId,
+          approvalId: passthroughPart.approvalId,
+          toolCallId: passthroughPart.toolCallId,
         } as UIMessageChunk;
       }
       if (
-        p.type === 'finish-step' ||
-        p.type === 'start-step' ||
-        p.type === 'tool-output-denied'
+        passthroughPart.type === 'finish-step' ||
+        passthroughPart.type === 'start-step' ||
+        passthroughPart.type === 'tool-output-denied'
       ) {
-        return p as UIMessageChunk;
+        return passthroughPart as UIMessageChunk;
       }
       return undefined;
     }