Support closure variables for serialized step functions

.changeset/bright-ducks-travel.md

@@ -0,0 +1,5 @@
+---
+"@workflow/core": patch
+---
+
+Support closure variables for serialized step functions

packages/core/e2e/e2e.test.ts

@@ -28,7 +28,7 @@
     withResolvers<string>();
   const dehydratedArgs = dehydrateWorkflowArguments(args, ops, runIdPromise);
 
   const res = await fetch(url, {
     method: 'POST',
     body: JSON.stringify(dehydratedArgs),
   });
@@ -88,10 +88,10 @@
       workflowFile: 'workflows/98_duplicate_case.ts',
       workflowFn: 'addTenWorkflow',
     },
   ])('addTenWorkflow', { timeout: 60_000 }, async (workflow) => {
     const run = await triggerWorkflow(workflow, [123]);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue).toBe(133);
 
     const { json } = await cliInspectJson(`runs ${run.runId} --withData`);
     expect(json).toMatchObject({
@@ -110,25 +110,25 @@
     ]);
   });
 
   test('promiseAllWorkflow', { timeout: 60_000 }, async () => {
     const run = await triggerWorkflow('promiseAllWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue).toBe('ABC');
   });
 
   test('promiseRaceWorkflow', { timeout: 60_000 }, async () => {
     const run = await triggerWorkflow('promiseRaceWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue).toBe('B');
   });
 
   test('promiseAnyWorkflow', { timeout: 60_000 }, async () => {
     const run = await triggerWorkflow('promiseAnyWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue).toBe('B');
   });
 
   test('readableStreamWorkflow', { timeout: 60_000 }, async () => {
     const run = await triggerWorkflow('readableStreamWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue).toBeInstanceOf(ReadableStream);
@@ -158,7 +158,7 @@
       method: 'POST',
       body: JSON.stringify({ token, data: { message: 'one' } }),
     });
     expect(res.status).toBe(200);
     let body = await res.json();
     expect(body.runId).toBe(run.runId);
 
@@ -194,7 +194,7 @@
     expect(returnValue).toBeInstanceOf(Array);
     expect(returnValue.length).toBe(3);
     expect(returnValue[0].message).toBe('one');
     expect(returnValue[0].customData).toBe(customData);
     expect(returnValue[0].done).toBeUndefined();
     expect(returnValue[1].message).toBe('two');
     expect(returnValue[1].customData).toBe(customData);
@@ -230,7 +230,7 @@
         body: JSON.stringify({ message: 'one' }),
       }
     );
     expect(res.status).toBe(202);
     const body = await res.text();
     expect(body).toBe('');
 
@@ -245,7 +245,7 @@
         body: JSON.stringify({ message: 'two' }),
       }
     );
     expect(res2.status).toBe(402);
     const body2 = await res2.text();
     expect(body2).toBe('Hello from static response!');
 
@@ -299,7 +299,7 @@
       `/.well-known/workflow/v1/webhook/${encodeURIComponent('invalid')}`,
       deploymentUrl
     );
     const res = await fetch(invalidWebhookUrl, {
       method: 'POST',
       body: JSON.stringify({}),
     });
@@ -308,17 +308,17 @@
     expect(body).toBe('');
   });
 
   test('sleepingWorkflow', { timeout: 60_000 }, async () => {
     const run = await triggerWorkflow('sleepingWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue.startTime).toBeLessThan(returnValue.endTime);
     expect(returnValue.endTime - returnValue.startTime).toBeGreaterThan(9999);
   });
 
   test('nullByteWorkflow', { timeout: 60_000 }, async () => {
     const run = await triggerWorkflow('nullByteWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue).toBe('null byte \0');
   });
 
   test('workflowAndStepMetadataWorkflow', { timeout: 60_000 }, async () => {
@@ -537,7 +537,7 @@
     const run = await triggerWorkflow('retryableAndFatalErrorWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);
     expect(returnValue.retryableResult.attempt).toEqual(2);
     expect(returnValue.retryableResult.duration).toBeGreaterThan(10_000);
     expect(returnValue.gotFatalError).toBe(true);
   });
 
@@ -661,7 +661,7 @@
 
       // Get first workflow result
       const run1Result = await getWorkflowReturnValue(run1.runId);
       expect(run1Result).toMatchObject({
         message: 'test-message-1',
         customData,
         hookCleanupTestData: 'workflow_completed',
@@ -707,7 +707,7 @@
   );
 
   test(
-    'stepFunctionPassingWorkflow - step function references can be passed as arguments',
+    'stepFunctionPassingWorkflow - step function references can be passed as arguments (without closure vars)',
     { timeout: 60_000 },
     async () => {
       // This workflow passes a step function reference to another step
@@ -738,6 +738,31 @@
     }
   );
 
+  test(
+    'stepFunctionWithClosureWorkflow - step function with closure variables passed as argument',
+    { timeout: 60_000 },
+    async () => {
+      // This workflow creates a nested step function with closure variables,
+      // then passes it to another step which invokes it.
+      // The closure variables should be serialized and preserved across the call.
+      const run = await triggerWorkflow('stepFunctionWithClosureWorkflow', []);
+      const returnValue = await getWorkflowReturnValue(run.runId);
+
+      // Expected: "Wrapped: Result: 21"
+      // - calculate(7) uses closure vars: prefix="Result: ", multiplier=3
+      // - 7 * 3 = 21, prefixed with "Result: " = "Result: 21"
+      // - stepThatCallsStepFn wraps it: "Wrapped: Result: 21"
+      expect(returnValue).toBe('Wrapped: Result: 21');
+
+      // Verify the run completed successfully
+      const { json: runData } = await cliInspectJson(
+        `runs ${run.runId} --withData`
+      );
+      expect(runData.status).toBe('completed');
+      expect(runData.output).toBe('Wrapped: Result: 21');
+    }
+  );
+
   test(
     'closureVariableWorkflow - nested step functions with closure variables',
     { timeout: 60_000 },

packages/core/src/serialization.test.ts

@@ -10,6 +10,7 @@ import {
   getCommonRevivers,
   getStreamType,
   getWorkflowReducers,
+  hydrateStepArguments,
   hydrateWorkflowArguments,
 } from './serialization.js';
 import { STREAM_NAME_SYMBOL } from './symbols.js';
@@ -844,7 +845,7 @@ describe('step function serialization', () => {
 
     // Get the reviver and test it directly
     const revivers = getCommonRevivers(vmGlobalThis);
-    const result = revivers.StepFunction(stepName);
+    const result = revivers.StepFunction({ stepId: stepName });
 
     expect(result).toBe(stepFn);
   });
@@ -854,7 +855,7 @@ describe('step function serialization', () => {
 
     let err: Error | undefined;
     try {
-      revivers.StepFunction('nonExistentStep');
+      revivers.StepFunction({ stepId: 'nonExistentStep' });
     } catch (err_) {
       err = err_ as Error;
     }
@@ -895,7 +896,88 @@ describe('step function serialization', () => {
     expect(dehydrated).toContain(42);
   });
 
-  it('should serialize step function to name through reducer', () => {
+  it('should dehydrate and hydrate step function with closure variables', async () => {
+    const stepName = 'step//workflows/test.ts//calculate';
+
+    // Create a step function that accesses closure variables
+    const { __private_getClosureVars } = await import('./private.js');
+    const { contextStorage } = await import('./step/context-storage.js');
+
+    const stepFn = async (x: number) => {
+      const { multiplier, prefix } = __private_getClosureVars();
+      const result = x * multiplier;
+      return `${prefix}${result}`;
+    };
+
+    // Register the step function
+    registerStepFunction(stepName, stepFn);
+
+    // Simulate what useStep() does - attach stepId and closure vars function
+    Object.defineProperty(stepFn, 'stepId', {
+      value: stepName,
+      writable: false,
+      enumerable: false,
+      configurable: false,
+    });
+
+    const closureVars = { multiplier: 3, prefix: 'Result: ' };
+    Object.defineProperty(stepFn, '__closureVarsFn', {
+      value: () => closureVars,
+      writable: false,
+      enumerable: false,
+      configurable: false,
+    });
+
+    // Serialize the step function with closure variables
+    const args = [stepFn, 7];
+    const dehydrated = dehydrateStepArguments(args, globalThis);
+
+    // Verify it serialized
+    expect(dehydrated).toBeDefined();
+    const serialized = JSON.stringify(dehydrated);
+    expect(serialized).toContain(stepName);
+    expect(serialized).toContain('multiplier');
+    expect(serialized).toContain('prefix');
+
+    // Now hydrate it back
+    const hydrated = hydrateStepArguments(
+      dehydrated,
+      [],
+      'test-run-123',
+      vmGlobalThis
+    );
+    expect(Array.isArray(hydrated)).toBe(true);
+    expect(hydrated).toHaveLength(2);
+
+    const hydratedStepFn = hydrated[0];
+    const hydratedArg = hydrated[1];
+
+    expect(typeof hydratedStepFn).toBe('function');
+    expect(hydratedArg).toBe(7);
+
+    // Invoke the hydrated step function within a context
+    const result = await contextStorage.run(
+      {
+        stepMetadata: {
+          stepId: 'test-step',
+          stepStartedAt: new Date(),
+          attempt: 1,
+        },
+        workflowMetadata: {
+          workflowRunId: 'test-run',
+          workflowStartedAt: new Date(),
+          url: 'http://localhost:3000',
+        },
+        ops: [],
+      },
+      () => hydratedStepFn(7)
+    );
+
+    // Verify the closure variables were accessible and used correctly
+    expect(result).toBe('Result: 21');
+  });
+
+  it('should serialize step function to object through reducer', () => {
     const stepName = 'step//workflows/test.ts//anotherStep';
     const stepFn = async () => 'result';
 
@@ -911,7 +993,7 @@ describe('step function serialization', () => {
     const reducer = getWorkflowReducers(globalThis).StepFunction;
     const result = reducer(stepFn);
 
-    // Should return the step name
-    expect(result).toBe(stepName);
+    // Should return object with stepId
+    expect(result).toEqual({ stepId: stepName });
   });
 });

packages/core/src/serialization.ts

@@ -2,6 +2,7 @@ import { WorkflowRuntimeError } from '@workflow/errors';
 import * as devalue from 'devalue';
 import { getStepFunction } from './private.js';
 import { getWorld } from './runtime/world.js';
+import { contextStorage } from './step/context-storage.js';
 import {
   BODY_INIT_SYMBOL,
   STREAM_NAME_SYMBOL,
@@ -181,7 +182,10 @@ export interface SerializableSpecial {
     redirected: boolean;
   };
   Set: any[];
-  StepFunction: string; // step function name/ID
+  StepFunction: {
+    stepId: string;
+    closureVars?: Record<string, any>;
+  };
   URL: string;
   URLSearchParams: string;
   Uint8Array: string; // base64 string
@@ -291,7 +295,18 @@ function getCommonReducers(global: Record<string, any> = globalThis) {
     StepFunction: (value) => {
       if (typeof value !== 'function') return false;
       const stepId = (value as any).stepId;
-      return typeof stepId === 'string' ? stepId : false;
+      if (typeof stepId !== 'string') return false;
+
+      // Check if the step function has closure variables
+      const closureVarsFn = (value as any).__closureVarsFn;
+      if (closureVarsFn && typeof closureVarsFn === 'function') {
+        // Invoke the closure variables function and serialize along with stepId
+        const closureVars = closureVarsFn();
+        return { stepId, closureVars };
+      }
+
+      // No closure variables - return object with just stepId
+      return { stepId };
     },
     URL: (value) => value instanceof global.URL && value.href,
     URLSearchParams: (value) => {
@@ -556,12 +571,56 @@ export function getCommonRevivers(global: Record<string, any> = globalThis) {
     RegExp: (value) => new global.RegExp(value.source, value.flags),
     Set: (value) => new global.Set(value),
     StepFunction: (value) => {
-      const stepFn = getStepFunction(value);
+      const stepId = value.stepId;
+      const closureVars = value.closureVars;
+
+      const stepFn = getStepFunction(stepId);
       if (!stepFn) {
         throw new Error(
-          `Step function "${value}" not found. Make sure the step function is registered.`
+          `Step function "${stepId}" not found. Make sure the step function is registered.`
         );
       }
+
+      // If closure variables were serialized, return a wrapper function
+      // that sets up AsyncLocalStorage context when invoked
+      if (closureVars) {
+        const wrappedStepFn = ((...args: any[]) => {
+          // Get the current context from AsyncLocalStorage
+          const currentContext = contextStorage.getStore();
+
+          if (!currentContext) {
+            throw new Error(
+              'Cannot call step function with closure variables outside step context'
+            );
+          }
+
+          // Create a new context with the closure variables merged in
+          const newContext = {
+            ...currentContext,
+            closureVars,
+          };
+
+          // Run the step function with the new context that includes closure vars
+          return contextStorage.run(newContext, () => stepFn(...args));
+        }) as any;
+
+        // Copy properties from original step function
+        Object.defineProperty(wrappedStepFn, 'name', {
+          value: stepFn.name,
+        });
+        Object.defineProperty(wrappedStepFn, 'stepId', {
+          value: stepId,
+          writable: false,
+          enumerable: false,
+          configurable: false,
+        });
+        if (stepFn.maxRetries !== undefined) {
+          wrappedStepFn.maxRetries = stepFn.maxRetries;
+        }
+
+        return wrappedStepFn;
+      }
+
       return stepFn;
     },
     URL: (value) => new global.URL(value),

packages/core/src/step.ts

@@ -150,6 +150,16 @@ export function createUseStep(ctx: WorkflowOrchestratorContext) {
       configurable: false,
     });
 
+    // Store the closure variables function for serialization
+    if (closureVarsFn) {
+      Object.defineProperty(stepFunction, '__closureVarsFn', {
+        value: closureVarsFn,
+        writable: false,
+        enumerable: false,
+        configurable: false,
+      });
+    }
+
     return stepFunction;
   };
 }