Add support for closure scope vars in step functions (#358)

### What changed?

- Enhanced the SWC plugin to detect and collect closure variables from nested step functions
- Modified the workflow runtime to pass closure variables to step functions during execution
- Updated the serialization/deserialization logic to handle the new closure variable format
- Added a mechanism to access closure variables within step functions
- Added tests to verify closure variable functionality in nested step functions
- Closure variables are stored on the step function execution's `AsyncLocalStorage`​ context

### Example

```typescript
export async function myWorkflow(baseValue: number) {
  'use workflow';
  const multiplier = 3;
  const prefix = 'Result: ';

  const calculate = async () => {
    'use step';
    const result = baseValue * multiplier;
    return `${prefix}${result}`;
  };

  return await calculate();
}
```

### Why make this change?

Previously, nested step functions couldn't access variables from their parent workflow scope, limiting their usefulness and requiring workarounds like passing all needed values as parameters. This change enables a more natural programming model where step functions can access variables from their enclosing scope, making workflow code more intuitive and reducing the need for explicit parameter passing.

.changeset/cruel-corners-feel.md

@@ -0,0 +1,8 @@
+---
+"@workflow/web-shared": patch
+"@workflow/swc-plugin": patch
+"@workflow/world": patch
+"@workflow/core": patch
+---
+
+Add support for closure scope vars in step functions

packages/core/e2e/e2e.test.ts

@@ -737,4 +737,18 @@ describe('e2e', () => {
       expect(stepCompletedEvents).toHaveLength(1);
     }
   );
+
+  test(
+    'closureVariableWorkflow - nested step functions with closure variables',
+    { timeout: 60_000 },
+    async () => {
+      // This workflow uses a nested step function that references closure variables
+      // from the parent workflow scope (multiplier, prefix, baseValue)
+      const run = await triggerWorkflow('closureVariableWorkflow', [7]);
+      const returnValue = await getWorkflowReturnValue(run.runId);
+
+      // Expected: baseValue (7) * multiplier (3) = 21, prefixed with "Result: "
+      expect(returnValue).toBe('Result: 21');
+    }
+  );
 });

packages/core/src/global.ts

@@ -5,6 +5,7 @@ export interface StepInvocationQueueItem {
   correlationId: string;
   stepName: string;
   args: Serializable[];
+  closureVars?: Record<string, Serializable>;
 }
 
 export interface HookInvocationQueueItem {

packages/core/src/private.ts

@@ -29,6 +29,12 @@ export function getStepFunction(stepId: string): StepFunction | undefined {
   return registeredSteps.get(stepId);
 }
 
+/**
+ * Get closure variables for the current step function
+ * @internal
+ */
+export { __private_getClosureVars } from './step/get-closure-vars.js';
+
 export interface WorkflowOrchestratorContext {
   globalThis: typeof globalThis;
   eventsConsumer: EventsConsumer;

packages/core/src/runtime.ts

@@ -389,8 +389,11 @@ export function workflowEntrypoint(workflowCode: string) {
                 if (queueItem.type === 'step') {
                   // Handle step operations
                   const ops: Promise<void>[] = [];
-                  const dehydratedArgs = dehydrateStepArguments(
-                    queueItem.args,
+                  const dehydratedInput = dehydrateStepArguments(
+                    {
+                      args: queueItem.args,
+                      closureVars: queueItem.closureVars,
+                    },
                     err.globalThis
                   );
 
@@ -398,7 +401,7 @@ export function workflowEntrypoint(workflowCode: string) {
                     const step = await world.steps.create(runId, {
                       stepId: queueItem.correlationId,
                       stepName: queueItem.stepName,
-                      input: dehydratedArgs as Serializable[],
+                      input: dehydratedInput as Serializable,
                     });
 
                     waitUntil(
@@ -678,9 +681,15 @@ export const stepEntrypoint =
                 `Step "${stepId}" has no "startedAt" timestamp`
               );
             }
-            // Hydrate the step input arguments
+            // Hydrate the step input arguments and closure variables
             const ops: Promise<void>[] = [];
-            const args = hydrateStepArguments(step.input, ops, workflowRunId);
+            const hydratedInput = hydrateStepArguments(
+              step.input,
+              ops,
+              workflowRunId
+            );
+
+            const args = hydratedInput.args;
 
             span?.setAttributes({
               ...Attribute.StepArgumentsCount(args.length),
@@ -703,8 +712,9 @@ export const stepEntrypoint =
                     : `http://localhost:${port ?? 3000}`,
                 },
                 ops,
+                closureVars: hydratedInput.closureVars,
               },
-              () => stepFn(...args)
+              () => stepFn.apply(null, args)
             );
 
             // NOTE: None of the code from this point is guaranteed to run

packages/core/src/step.test.ts

@@ -199,4 +199,75 @@ describe('createUseStep', () => {
     await myStepFunction();
     expect(ctx.onWorkflowError).not.toHaveBeenCalled();
   });
+
+  it('should capture closure variables when provided', async () => {
+    const ctx = setupWorkflowContext([
+      {
+        eventId: 'evnt_0',
+        runId: 'wrun_123',
+        eventType: 'step_completed',
+        correlationId: 'step_01K11TFZ62YS0YYFDQ3E8B9YCV',
+        eventData: {
+          result: ['Result: 42'],
+        },
+        createdAt: new Date(),
+      },
+    ]);
+
+    const useStep = createUseStep(ctx);
+    const count = 42;
+    const prefix = 'Result: ';
+
+    // Create step with closure variables function
+    const calculate = useStep('calculate', () => ({ count, prefix }));
+
+    // Call the step
+    const result = await calculate();
+
+    // Verify result
+    expect(result).toBe('Result: 42');
+
+    // Verify closure variables were added to invocation queue
+    expect(ctx.invocationsQueue).toHaveLength(1);
+    expect(ctx.invocationsQueue[0]).toMatchObject({
+      type: 'step',
+      stepName: 'calculate',
+      args: [],
+      closureVars: { count: 42, prefix: 'Result: ' },
+    });
+  });
+
+  it('should handle empty closure variables', async () => {
+    const ctx = setupWorkflowContext([
+      {
+        eventId: 'evnt_0',
+        runId: 'wrun_123',
+        eventType: 'step_completed',
+        correlationId: 'step_01K11TFZ62YS0YYFDQ3E8B9YCV',
+        eventData: {
+          result: [5],
+        },
+        createdAt: new Date(),
+      },
+    ]);
+
+    const useStep = createUseStep(ctx);
+
+    // Create step without closure variables
+    const add = useStep('add');
+
+    // Call the step
+    const result = await add(2, 3);
+
+    // Verify result
+    expect(result).toBe(5);
+
+    // Verify empty closure variables were added to invocation queue
+    expect(ctx.invocationsQueue).toHaveLength(1);
+    expect(ctx.invocationsQueue[0]).toMatchObject({
+      type: 'step',
+      stepName: 'add',
+      args: [2, 3],
+    });
+  });
 });

packages/core/src/step.ts

@@ -1,7 +1,7 @@
 import { FatalError, WorkflowRuntimeError } from '@workflow/errors';
 import { withResolvers } from '@workflow/utils';
 import { EventConsumerResult } from './events-consumer.js';
-import { WorkflowSuspension } from './global.js';
+import { type StepInvocationQueueItem, WorkflowSuspension } from './global.js';
 import { stepLogger } from './logger.js';
 import type { WorkflowOrchestratorContext } from './private.js';
 import type { Serializable } from './schemas.js';
@@ -9,18 +9,28 @@ import { hydrateStepReturnValue } from './serialization.js';
 
 export function createUseStep(ctx: WorkflowOrchestratorContext) {
   return function useStep<Args extends Serializable[], Result>(
-    stepName: string
+    stepName: string,
+    closureVarsFn?: () => Record<string, Serializable>
   ) {
     const stepFunction = (...args: Args): Promise<Result> => {
       const { promise, resolve, reject } = withResolvers<Result>();
 
       const correlationId = `step_${ctx.generateUlid()}`;
-      ctx.invocationsQueue.push({
+
+      const queueItem: StepInvocationQueueItem = {
         type: 'step',
         correlationId,
         stepName,
         args,
-      });
+      };
+
+      // Invoke the closure variables function to get the closure scope
+      const closureVars = closureVarsFn?.();
+      if (closureVars) {
+        queueItem.closureVars = closureVars;
+      }
+
+      ctx.invocationsQueue.push(queueItem);
 
       // Track whether we've already seen a "step_started" event for this step.
       // This is important because after a retryable failure, the step moves back to

packages/core/src/step/context-storage.ts

@@ -6,4 +6,5 @@ export const contextStorage = /* @__PURE__ */ new AsyncLocalStorage<{
   stepMetadata: StepMetadata;
   workflowMetadata: WorkflowMetadata;
   ops: Promise<void>[];
+  closureVars?: Record<string, any>;
 }>();

packages/core/src/step/get-closure-vars.ts

@@ -0,0 +1,18 @@
+import { contextStorage } from './context-storage.js';
+
+/**
+ * Returns the closure variables for the current step function.
+ * This is an internal function used by the SWC transform to access
+ * variables from the parent workflow scope.
+ *
+ * @internal
+ */
+export function __private_getClosureVars(): Record<string, any> {
+  const ctx = contextStorage.getStore();
+  if (!ctx) {
+    throw new Error(
+      'Closure variables can only be accessed inside a step function'
+    );
+  }
+  return ctx.closureVars || {};
+}

packages/core/src/workflow.test.ts

@@ -2304,4 +2304,96 @@ describe('runWorkflow', () => {
       );
     });
   });
+
+  describe('closure variables', () => {
+    it('should serialize and deserialize closure variables for nested step functions', async () => {
+      let error: Error | undefined;
+      try {
+        const ops: Promise<any>[] = [];
+        const workflowRun: WorkflowRun = {
+          runId: 'test-run-123',
+          workflowName: 'workflow',
+          status: 'running',
+          input: dehydrateWorkflowArguments([], ops),
+          createdAt: new Date('2024-01-01T00:00:00.000Z'),
+          updatedAt: new Date('2024-01-01T00:00:00.000Z'),
+          startedAt: new Date('2024-01-01T00:00:00.000Z'),
+          deploymentId: 'test-deployment',
+        };
+
+        const events: Event[] = [];
+
+        await runWorkflow(
+          `const useStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")];
+          async function workflow() {
+            const multiplier = 3;
+            const prefix = 'Result: ';
+            const calculate = useStep('step//input.js//_anonymousStep0', () => ({ multiplier, prefix }));
+            const result = await calculate(7);
+            return result;
+          }${getWorkflowTransformCode('workflow')}`,
+          workflowRun,
+          events
+        );
+      } catch (err) {
+        error = err as Error;
+      }
+
+      // Should suspend to create the step
+      assert(error);
+      expect(error.name).toEqual('WorkflowSuspension');
+      expect((error as WorkflowSuspension).steps).toHaveLength(1);
+
+      const step = (error as WorkflowSuspension).steps[0];
+      expect(step).toMatchObject({
+        type: 'step',
+        stepName: 'step//input.js//_anonymousStep0',
+        args: [7],
+        closureVars: { multiplier: 3, prefix: 'Result: ' },
+      });
+    });
+
+    it('should handle step functions without closure variables', async () => {
+      let error: Error | undefined;
+      try {
+        const ops: Promise<any>[] = [];
+        const workflowRun: WorkflowRun = {
+          runId: 'test-run-123',
+          workflowName: 'workflow',
+          status: 'running',
+          input: dehydrateWorkflowArguments([], ops),
+          createdAt: new Date('2024-01-01T00:00:00.000Z'),
+          updatedAt: new Date('2024-01-01T00:00:00.000Z'),
+          startedAt: new Date('2024-01-01T00:00:00.000Z'),
+          deploymentId: 'test-deployment',
+        };
+
+        const events: Event[] = [];
+
+        await runWorkflow(
+          `const add = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("add");
+          async function workflow() {
+            const result = await add(5, 10);
+            return result;
+          }${getWorkflowTransformCode('workflow')}`,
+          workflowRun,
+          events
+        );
+      } catch (err) {
+        error = err as Error;
+      }
+
+      // Should suspend to create the step
+      assert(error);
+      expect(error.name).toEqual('WorkflowSuspension');
+      expect((error as WorkflowSuspension).steps).toHaveLength(1);
+
+      const step = (error as WorkflowSuspension).steps[0];
+      expect(step).toMatchObject({
+        type: 'step',
+        stepName: 'add',
+        args: [5, 10],
+      });
+    });
+  });
 });

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-step.js

@@ -4,15 +4,9 @@ import { registerStepFunction } from "workflow/internal/private";
 async function step(a, b) {
     return a + b;
 }
-async function arrowStep(x, y) {
-    return x * y;
-}
-async function letArrowStep(x, y) {
-    return x - y;
-}
-async function varArrowStep(x, y) {
-    return x / y;
-}
+var arrowStep = async (x, y)=>x * y;
+var letArrowStep = async (x, y)=>x - y;
+var varArrowStep = async (x, y)=>x / y;
 var helpers$objectStep = async (x, y)=>{
     return x + y + 10;
 };

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/input.js

@@ -0,0 +1,30 @@
+import { DurableAgent } from '@workflow/ai/agent';
+import { gateway } from 'ai';
+
+export async function wflow() {
+  'use workflow';
+  let count = 42;
+
+  async function namedStepWithClosureVars() {
+    'use step';
+    console.log('count', count);
+  }
+
+  const agent = new DurableAgent({
+    arrowFunctionWithClosureVars: async () => {
+      'use step';
+      console.log('count', count);
+      return gateway('openai/gpt-5');
+    },
+
+    namedFunctionWithClosureVars: async function() {
+      'use step';
+      console.log('count', count);
+    },
+
+    async methodWithClosureVars() {
+      'use step';
+      console.log('count', count);
+    },
+  });
+}

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-client.js

@@ -0,0 +1,5 @@
+/**__internal_workflows{"workflows":{"input.js":{"wflow":{"workflowId":"workflow//input.js//wflow"}}}}*/;
+export async function wflow() {
+    throw new Error("You attempted to execute workflow wflow function directly. To start a workflow, use start(wflow) from workflow/api");
+}
+wflow.workflowId = "workflow//input.js//wflow";

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-step.js

@@ -0,0 +1,33 @@
+import { __private_getClosureVars, registerStepFunction } from "workflow/internal/private";
+import { DurableAgent } from '@workflow/ai/agent';
+import { gateway } from 'ai';
+/**__internal_workflows{"workflows":{"input.js":{"wflow":{"workflowId":"workflow//input.js//wflow"}}},"steps":{"input.js":{"_anonymousStep0":{"stepId":"step//input.js//_anonymousStep0"},"_anonymousStep1":{"stepId":"step//input.js//_anonymousStep1"},"_anonymousStep2":{"stepId":"step//input.js//_anonymousStep2"},"namedStepWithClosureVars":{"stepId":"step//input.js//namedStepWithClosureVars"}}}}*/;
+async function namedStepWithClosureVars() {
+    const { count } = __private_getClosureVars();
+    console.log('count', count);
+}
+var _anonymousStep0 = async ()=>{
+    const { count } = __private_getClosureVars();
+    console.log('count', count);
+    return gateway('openai/gpt-5');
+};
+async function _anonymousStep1() {
+    const { count } = __private_getClosureVars();
+    console.log('count', count);
+}
+async function _anonymousStep2() {
+    const { count } = __private_getClosureVars();
+    console.log('count', count);
+}
+export async function wflow() {
+    let count = 42;
+    const agent = new DurableAgent({
+        arrowFunctionWithClosureVars: _anonymousStep0,
+        namedFunctionWithClosureVars: _anonymousStep1,
+        methodWithClosureVars: _anonymousStep2
+    });
+}
+registerStepFunction("step//input.js//namedStepWithClosureVars", namedStepWithClosureVars);
+registerStepFunction("step//input.js//_anonymousStep0", _anonymousStep0);
+registerStepFunction("step//input.js//_anonymousStep1", _anonymousStep1);
+registerStepFunction("step//input.js//_anonymousStep2", _anonymousStep2);

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-workflow.js

@@ -0,0 +1,20 @@
+import { DurableAgent } from '@workflow/ai/agent';
+/**__internal_workflows{"workflows":{"input.js":{"wflow":{"workflowId":"workflow//input.js//wflow"}}},"steps":{"input.js":{"_anonymousStep0":{"stepId":"step//input.js//_anonymousStep0"},"_anonymousStep1":{"stepId":"step//input.js//_anonymousStep1"},"_anonymousStep2":{"stepId":"step//input.js//_anonymousStep2"},"namedStepWithClosureVars":{"stepId":"step//input.js//namedStepWithClosureVars"}}}}*/;
+export async function wflow() {
+    let count = 42;
+    var namedStepWithClosureVars = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//namedStepWithClosureVars", ()=>({
+            count
+        }));
+    const agent = new DurableAgent({
+        arrowFunctionWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep0", ()=>({
+                count
+            })),
+        namedFunctionWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep1", ()=>({
+                count
+            })),
+        methodWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep2", ()=>({
+                count
+            }))
+    });
+}
+wflow.workflowId = "workflow//input.js//wflow";

packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-step.js

@@ -3,12 +3,8 @@ import { DurableAgent } from '@workflow/ai/agent';
 import { gateway, tool } from 'ai';
 import * as z from 'zod';
 /**__internal_workflows{"workflows":{"input.js":{"test":{"workflowId":"workflow//input.js//test"}}},"steps":{"input.js":{"_anonymousStep0":{"stepId":"step//input.js//_anonymousStep0"},"_anonymousStep1":{"stepId":"step//input.js//_anonymousStep1"}}}}*/;
-async function _anonymousStep0() {
-    return gateway('openai/gpt-5');
-}
-async function _anonymousStep1({ location }) {
-    return `Weather in ${location}: Sunny, 72°F`;
-}
+var _anonymousStep0 = async ()=>gateway('openai/gpt-5');
+var _anonymousStep1 = async ({ location })=>`Weather in ${location}: Sunny, 72°F`;
 export async function test() {
     'use workflow';
     const agent = new DurableAgent({

packages/web-shared/src/sidebar/attribute-panel.tsx

@@ -105,6 +105,36 @@ const attributeToDisplayFn: Record<
   // Resolved attributes, won't actually use this function
   metadata: JsonBlock,
   input: (value: unknown) => {
+    // Check if input has args + closure vars structure
+    if (value && typeof value === 'object' && 'args' in value) {
+      const { args, closureVars } = value as {
+        args: unknown[];
+        closureVars?: Record<string, unknown>;
+      };
+      const argCount = Array.isArray(args) ? args.length : 0;
+      const hasClosureVars = closureVars && Object.keys(closureVars).length > 0;
+
+      return (
+        <>
+          <DetailCard summary={`Input (${argCount} arguments)`}>
+            {Array.isArray(args)
+              ? args.map((v, i) => (
+                  <div className="mt-2" key={i}>
+                    {JsonBlock(v)}
+                  </div>
+                ))
+              : JsonBlock(args)}
+          </DetailCard>
+          {hasClosureVars && (
+            <DetailCard summary="Closure Variables">
+              {JsonBlock(closureVars)}
+            </DetailCard>
+          )}
+        </>
+      );
+    }
+
+    // Fallback: treat as plain array or object
     const argCount = Array.isArray(value) ? value.length : 0;
     return (
       <DetailCard summary={`Input (${argCount} arguments)`}>

packages/world/src/steps.ts

@@ -41,7 +41,7 @@ export type Step = z.infer<typeof StepSchema>;
 export interface CreateStepRequest {
   stepId: string;
   stepName: string;
-  input: SerializedData[];
+  input: SerializedData;
 }
 
 export interface UpdateStepRequest {

workbench/example/workflows/99_e2e.ts

@@ -512,3 +512,21 @@ async function doubleNumber(x: number) {
   'use step';
   return x * 2;
 }
+
+//////////////////////////////////////////////////////////
+
+export async function closureVariableWorkflow(baseValue: number) {
+  'use workflow';
+  let multiplier = 3;
+  const prefix = 'Result: ';
+
+  // Nested step function that uses closure variables
+  const calculate = async () => {
+    'use step';
+    const result = baseValue * multiplier;
+    return `${prefix}${result}`;
+  };
+
+  const output = await calculate();
+  return output;
+}