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.

Author: TooTallNate

.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,16 +389,19 @@ 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
                   );
 
                   try {
                     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,26 +1,36 @@
 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';
 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],
+      });
+    });
+  });
 });