@ai-sdk/workflow 1.0.0-beta.103 → 1.0.0-beta.105

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ai-sdk/workflow",
-  "version": "1.0.0-beta.103",
+  "version": "1.0.0-beta.105",
   "description": "WorkflowAgent for building AI agents with AI SDK",
   "license": "Apache-2.0",
   "main": "./dist/index.js",
@@ -29,7 +29,7 @@
     "ajv": "^8.20.0",
     "@ai-sdk/provider": "4.0.0-beta.20",
     "@ai-sdk/provider-utils": "5.0.0-beta.50",
-    "ai": "7.0.0-beta.184"
+    "ai": "7.0.0-beta.185"
   },
   "devDependencies": {
     "@types/node": "22.19.19",

package/src/stream-text-iterator.ts

@@ -7,6 +7,7 @@ import type { Context } from '@ai-sdk/provider-utils';
 import {
   experimental_filterActiveTools as filterActiveTools,
   type Experimental_LanguageModelStreamPart as ModelCallStreamPart,
+  type Experimental_SandboxSession as SandboxSession,
   type LanguageModel,
   type ModelMessage,
   type StepResult,
@@ -52,6 +53,8 @@ export interface StreamTextIteratorYieldValue {
   toolsContext?: Record<string, Context | undefined>;
   /** Provider-executed tool results (keyed by tool call ID) */
   providerExecutedToolResults?: Map<string, ProviderExecutedToolResult>;
+  /** The sandbox selected for the current step. */
+  experimental_sandbox?: SandboxSession;
 }
 
 // This runs in the workflow context
@@ -74,6 +77,7 @@ export async function* streamTextIterator({
   includeRawChunks = false,
   repairToolCall,
   responseFormat,
+  experimental_sandbox: sandbox,
 }: {
   prompt: LanguageModelV4Prompt;
   tools: ToolSet;
@@ -94,6 +98,7 @@ export async function* streamTextIterator({
   includeRawChunks?: boolean;
   repairToolCall?: ToolCallRepairFunction<ToolSet>;
   responseFormat?: LanguageModelV4CallOptions['responseFormat'];
+  experimental_sandbox?: SandboxSession;
 }): AsyncGenerator<
   StreamTextIteratorYieldValue,
   LanguageModelV4Prompt,
@@ -135,6 +140,8 @@ export async function* streamTextIterator({
       break;
     }
 
+    let stepSandbox = sandbox;
+
     // Call prepareStep callback before each step if provided
     if (prepareStep) {
       const prepareResult = await prepareStep({
@@ -144,8 +151,11 @@ export async function* streamTextIterator({
         messages: conversationPrompt,
         runtimeContext: currentRuntimeContext,
         toolsContext: currentToolsContext as never,
+        experimental_sandbox: sandbox,
       });
 
+      stepSandbox = prepareResult?.experimental_sandbox ?? sandbox;
+
       // Apply any overrides from prepareStep
       if (prepareResult?.model !== undefined) {
         currentModel = prepareResult.model;
@@ -398,6 +408,7 @@ export async function* streamTextIterator({
           step,
           runtimeContext: currentRuntimeContext,
           toolsContext: currentToolsContext,
+          experimental_sandbox: stepSandbox,
           providerExecutedToolResults,
         };
 
@@ -473,6 +484,7 @@ export async function* streamTextIterator({
       step: lastStep,
       runtimeContext: currentRuntimeContext,
       toolsContext: currentToolsContext,
+      experimental_sandbox: sandbox,
     };
   }
 

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

@@ -4,6 +4,7 @@
 import { tool } from 'ai';
 import { WorkflowAgent } from '../workflow-agent.js';
 import { mockTextModel, mockSequenceModel } from '../providers/mock.js';
+import { createTestSandbox } from './test-sandbox.js';
 import { FatalError, getWritable } from 'workflow';
 import { z } from 'zod';
 
@@ -585,3 +586,88 @@ export async function agentRuntimeAndToolsContextE2e() {
     onFinishToolsContext,
   };
 }
+
+export async function agentSandboxE2e() {
+  'use workflow';
+
+  let constructorSandboxRanCommand = 'not-run';
+  let stepSandboxRanCommand = 'not-run';
+  let firstPrepareStepSawConstructorSandbox = false;
+  let secondPrepareStepSawConstructorSandbox = false;
+  let prepareStepSawStepSandbox = false;
+
+  // A live sandbox session passed through `experimental_sandbox`. The tool
+  // `execute` is inline (not a `'use step'`) so the handle never crosses a
+  // step boundary — matching the single-process use of `experimental_sandbox`.
+  const constructorSandbox = createTestSandbox({
+    run: async ({ command }) => {
+      constructorSandboxRanCommand = command;
+      return {
+        exitCode: 0,
+        stdout: `constructor ran: ${command}`,
+        stderr: '',
+      };
+    },
+  });
+  const stepSandbox = createTestSandbox({
+    run: async ({ command }) => {
+      stepSandboxRanCommand = command;
+      return { exitCode: 0, stdout: `ran: ${command}`, stderr: '' };
+    },
+  });
+
+  const agent = new WorkflowAgent({
+    model: mockSequenceModel([
+      {
+        type: 'tool-call',
+        toolName: 'runShell',
+        input: JSON.stringify({ command: 'echo hello' }),
+      },
+      { type: 'text', text: 'Command executed.' },
+    ]),
+    tools: {
+      runShell: tool({
+        description: 'Run a shell command in the sandbox.',
+        inputSchema: z.object({ command: z.string() }),
+        execute: async ({ command }, { experimental_sandbox }) => {
+          if (experimental_sandbox == null) {
+            throw new Error('Sandbox is not available');
+          }
+          return experimental_sandbox.run({ command });
+        },
+      }),
+    },
+    instructions: 'You run shell commands in a sandbox.',
+    experimental_sandbox: constructorSandbox,
+    prepareStep: ({ stepNumber, experimental_sandbox }) => {
+      if (stepNumber === 0) {
+        firstPrepareStepSawConstructorSandbox =
+          experimental_sandbox === constructorSandbox;
+        return { experimental_sandbox: stepSandbox };
+      }
+
+      if (experimental_sandbox === stepSandbox) {
+        prepareStepSawStepSandbox = true;
+      }
+      secondPrepareStepSawConstructorSandbox =
+        experimental_sandbox === constructorSandbox;
+      return {};
+    },
+  });
+
+  const result = await agent.stream({
+    messages: [{ role: 'user', content: 'Run echo hello' }],
+    writable: getWritable(),
+  });
+
+  return {
+    stepCount: result.steps.length,
+    lastStepText: result.steps[result.steps.length - 1]?.text,
+    toolResults: result.toolResults,
+    constructorSandboxRanCommand,
+    stepSandboxRanCommand,
+    firstPrepareStepSawConstructorSandbox,
+    secondPrepareStepSawConstructorSandbox,
+    prepareStepSawStepSandbox,
+  };
+}

package/src/test/test-sandbox.ts

@@ -0,0 +1,26 @@
+import type { Experimental_SandboxSession as SandboxSession } from 'ai';
+
+/**
+ * Build a minimal {@link SandboxSession} for tests. Only `run` returns a usable
+ * result by default; every other method throws so tests stay explicit about
+ * what they exercise. Pass `overrides` to customize individual methods.
+ */
+export function createTestSandbox(
+  overrides: Partial<SandboxSession> = {},
+): SandboxSession {
+  const notImplemented = () => {
+    throw new Error('not implemented in test sandbox');
+  };
+  return {
+    description: 'test sandbox',
+    readFile: notImplemented,
+    readBinaryFile: notImplemented,
+    readTextFile: notImplemented,
+    writeFile: notImplemented,
+    writeBinaryFile: notImplemented,
+    writeTextFile: notImplemented,
+    spawn: notImplemented,
+    run: async () => ({ exitCode: 0, stdout: '', stderr: '' }),
+    ...overrides,
+  };
+}

package/src/workflow-agent.ts

@@ -8,6 +8,7 @@ import type {
 import {
   getErrorMessage,
   validateTypes,
+  withUserAgentSuffix,
   type Context,
   type HasRequiredKey,
   type InferToolSetContext,
@@ -32,6 +33,7 @@ import {
   type Prompt,
   type TelemetryOptions as CoreTelemetryOptions,
   type Instructions,
+  type Experimental_SandboxSession as SandboxSession,
 } from 'ai';
 import {
   createRestrictedTelemetryDispatcher,
@@ -278,6 +280,11 @@ export interface PrepareStepInfo<
    * `prepareStep` to update it for the current and subsequent steps.
    */
   toolsContext: InferToolSetContext<TTools>;
+
+  /**
+   * The sandbox environment that the step is operating in.
+   */
+  experimental_sandbox?: SandboxSession;
 }
 
 /**
@@ -326,6 +333,11 @@ export interface PrepareStepResult<
    * Returning a value replaces the agent's tools context.
    */
   toolsContext?: InferToolSetContext<TTools>;
+
+  /**
+   * Override the sandbox environment for this step.
+   */
+  experimental_sandbox?: SandboxSession;
 }
 
 /**
@@ -494,6 +506,14 @@ export type WorkflowAgentOptions<
      */
     experimental_download?: DownloadFunction;
 
+    /**
+     * Default sandbox environment passed through to tool execution as
+     * `experimental_sandbox`.
+     *
+     * Per-stream `experimental_sandbox` values passed to `stream()` override this default.
+     */
+    experimental_sandbox?: SandboxSession;
+
     /**
      * Default callback function called before each step in the agent loop.
      * Use this to modify settings, manage context, or inject messages dynamically
@@ -933,6 +953,12 @@ export type WorkflowAgentStreamOptions<
      */
     experimental_download?: DownloadFunction;
 
+    /**
+     * Sandbox environment passed through to tool execution as
+     * `experimental_sandbox`. Overrides the constructor-level value if provided.
+     */
+    experimental_sandbox?: SandboxSession;
+
     /**
      * Callback function to be called after each step completes.
      */
@@ -1177,6 +1203,7 @@ export class WorkflowAgent<
   private output?: OutputSpecification<any, any>;
   private experimentalRepairToolCall?: ToolCallRepairFunction<TBaseTools>;
   private experimentalDownload?: DownloadFunction;
+  private experimentalSandbox?: SandboxSession;
   private prepareStep?: PrepareStepCallback<TBaseTools, TRuntimeContext>;
   private allowSystemInMessages: boolean;
   private constructorOnStepEnd?: WorkflowAgentOnStepEndCallback<
@@ -1214,6 +1241,7 @@ export class WorkflowAgent<
     this.output = options.output;
     this.experimentalRepairToolCall = options.experimental_repairToolCall;
     this.experimentalDownload = options.experimental_download;
+    this.experimentalSandbox = options.experimental_sandbox;
     this.prepareStep = options.prepareStep;
     this.constructorOnStepEnd = options.onStepEnd ?? options.onStepFinish;
     const { onFinish, onEnd = onFinish } = options;
@@ -1360,6 +1388,7 @@ export class WorkflowAgent<
         : { messages: effectiveMessages! }),
     } as Prompt);
     const download = options.experimental_download ?? this.experimentalDownload;
+    const sandbox = options.experimental_sandbox ?? this.experimentalSandbox;
 
     // Process tool approval responses before starting the agent loop.
     // This mirrors how stream-text.ts handles tool-approval-response parts:
@@ -1517,6 +1546,7 @@ export class WorkflowAgent<
                 toolCallId: approval.toolCallId,
                 messages: [],
                 context: resolvedContext,
+                experimental_sandbox: sandbox,
               });
             const toolResult =
               telemetryDispatcher.executeTool != null
@@ -1719,6 +1749,14 @@ export class WorkflowAgent<
       }),
     };
 
+    // tag the outgoing request so usage can be attributed to WorkflowAgent.
+    // chains with the `ai/<version>` and `ai-sdk/<provider>/<version>` suffixes
+    // added downstream by the model run and the provider.
+    mergedGenerationSettings.headers = withUserAgentSuffix(
+      mergedGenerationSettings.headers ?? {},
+      'ai-sdk-agent/workflow',
+    );
+
     // Merge constructor + stream callbacks (constructor first, then stream)
     const mergedOnStepEnd = mergeCallbacks(
       this.constructorOnStepEnd as
@@ -1821,6 +1859,7 @@ export class WorkflowAgent<
       messages: LanguageModelV4Prompt,
       perToolContexts: Record<string, Context | undefined>,
       currentStepNumber: number = 0,
+      stepSandbox?: SandboxSession,
     ): Promise<WorkflowToolExecutionResult> => {
       const toolCallEvent: ToolCall = {
         type: 'tool-call',
@@ -1862,7 +1901,14 @@ export class WorkflowAgent<
       let result: WorkflowToolExecutionResult;
       try {
         const execute = () =>
-          executeTool(toolCall, tools, messages, resolvedContext, download);
+          executeTool(
+            toolCall,
+            tools,
+            messages,
+            resolvedContext,
+            download,
+            stepSandbox,
+          );
         result =
           telemetryDispatcher.executeTool != null
             ? await telemetryDispatcher.executeTool({
@@ -2034,6 +2080,7 @@ export class WorkflowAgent<
         | ToolCallRepairFunction<ToolSet>
         | undefined,
       responseFormat: await (options.output ?? this.output)?.responseFormat,
+      experimental_sandbox: sandbox,
     });
 
     // Track the final conversation messages from the iterator
@@ -2059,8 +2106,10 @@ export class WorkflowAgent<
           step,
           runtimeContext: yieldedRuntimeContext,
           toolsContext: yieldedToolsContext,
+          experimental_sandbox: stepSandbox,
           providerExecutedToolResults,
         } = result.value;
+        const toolExecutionSandbox = stepSandbox ?? sandbox;
         // Capture current step number before pushing (0-based)
         const currentStepNumber = steps.length;
         if (step) {
@@ -2141,6 +2190,7 @@ export class WorkflowAgent<
                     iterMessages,
                     toolsContext,
                     currentStepNumber,
+                    toolExecutionSandbox,
                   ),
               ),
             );
@@ -2280,6 +2330,7 @@ export class WorkflowAgent<
                   iterMessages,
                   toolsContext,
                   currentStepNumber,
+                  toolExecutionSandbox,
                 ),
             ),
           );
@@ -2757,6 +2808,7 @@ async function executeTool(
   messages: LanguageModelV4Prompt,
   context?: unknown,
   download?: DownloadFunction,
+  sandbox?: SandboxSession,
 ): Promise<WorkflowToolExecutionResult> {
   const tool = tools[toolCall.toolName];
   if (!tool) throw new Error(`Tool "${toolCall.toolName}" not found`);
@@ -2783,6 +2835,7 @@ async function executeTool(
       messages,
       // Pass per-tool context to the tool (resolved from `toolsContext`)
       context,
+      experimental_sandbox: sandbox,
     });
   } catch (error) {
     // Convert tool errors to error-text results sent back to the model,