WIP closure scope vars in step functions

Author: TooTallNate

.changeset/cruel-corners-feel.md

@@ -0,0 +1,7 @@
+---
+"@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/observability.ts

@@ -49,18 +49,20 @@ const hydrateStepIO = <
 >(
   step: T
 ): T => {
+  const hydratedInput = step.input
+    ? hydrateStepArguments(
+        step.input,
+        [],
+        step.runId as string,
+        globalThis,
+        streamPrintRevivers
+      )
+    : step.input;
+
   return {
     ...step,
-    input:
-      step.input && Array.isArray(step.input) && step.input.length
-        ? hydrateStepArguments(
-            step.input,
-            [],
-            step.runId as string,
-            globalThis,
-            streamPrintRevivers
-          )
-        : step.input,
+    input: 'args' in hydratedInput ? hydratedInput.args : hydratedInput,
+    // TODO: should we also render the closure vars for o11y?
     output: step.output
       ? hydrateStepReturnValue(step.output, globalThis, streamPrintRevivers)
       : step.output,

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 || {};
+}