Support passing step function references across serialization boundaries

packages/core/e2e/e2e.test.ts

@@ -214,7 +214,7 @@
         body: JSON.stringify({ message: 'one' }),
       }
     );
     expect(res.status).toBe(202);
     const body = await res.text();
     expect(body).toBe('');
 
@@ -287,7 +287,7 @@
       method: 'POST',
       body: JSON.stringify({}),
     });
     expect(res.status).toBe(404);
     const body = await res.text();
     expect(body).toBe('');
   });
@@ -598,4 +598,25 @@
       expect(runData.error).toContain('Error from imported helper module');
     }
   );
+
+  test(
+    'stepFunctionPassingWorkflow - step function references can be passed as arguments',
+    { timeout: 60_000 },
+    async () => {
+      // This workflow passes a step function reference to another step
+      // The receiving step calls the passed function and returns the result
+      const run = await triggerWorkflow('stepFunctionPassingWorkflow', []);
+      const returnValue = await getWorkflowReturnValue(run.runId);
+
+      // doubleNumber(10) = 20, then multiply by 2 = 40
+      expect(returnValue).toBe(40);
+
+      // Verify the run completed successfully
+      const { json: runData } = await cliInspectJson(
+        `runs ${run.runId} --withData`
+      );
+      expect(runData.status).toBe('completed');
+      expect(runData.output).toBe(40);
+    }
+  );
 });

packages/core/src/runtime/world.ts

@@ -1,10 +1,10 @@
 import { createRequire } from 'node:module';
-import Path from 'node:path';
+import { join } from 'node:path';
 import type { World } from '@workflow/world';
 import { createEmbeddedWorld } from '@workflow/world-local';
 import { createVercelWorld } from '@workflow/world-vercel';
 
-const require = createRequire(Path.join(process.cwd(), 'index.js'));
+const require = createRequire(join(process.cwd(), 'index.js'));
 
 let worldCache: World | undefined;
 let stubbedWorldCache: World | undefined;

packages/core/src/schemas.ts

@@ -63,4 +63,5 @@ export type Serializable =
   | Uint8ClampedArray
   | Uint16Array
   | Uint32Array
-  | WritableStream<Uint8Array>;
+  | WritableStream<Uint8Array>
+  | ((...args: Serializable[]) => Promise<Serializable>); // Step function

packages/core/src/serialization.test.ts

@@ -6,9 +6,16 @@ import {
   dehydrateStepReturnValue,
   dehydrateWorkflowArguments,
   dehydrateWorkflowReturnValue,
+  getCommonRevivers,
   getStreamType,
+  hydrateStepArguments,
   hydrateWorkflowArguments,
 } from './serialization.js';
+import {
+  getStepFunction,
+  registerStepFunction,
+  STEP_FUNCTION_NAME_SYMBOL,
+} from './private.js';
 import { STREAM_NAME_SYMBOL } from './symbols.js';
 import { createContext } from './vm/index.js';
 
@@ -783,3 +790,79 @@ describe('step return value', () => {
     );
   });
 });
+
+describe('step function serialization', () => {
+  const { globalThis: vmGlobalThis } = createContext({
+    seed: 'test',
+    fixedTimestamp: 1714857600000,
+  });
+
+  it('should detect step function by checking for STEP_FUNCTION_NAME_SYMBOL', () => {
+    const stepName = 'myStep';
+    const stepFn = async (x: number) => x * 2;
+
+    // Attach the symbol like useStep() does
+    Object.defineProperty(stepFn, STEP_FUNCTION_NAME_SYMBOL, {
+      value: stepName,
+      writable: false,
+      enumerable: false,
+      configurable: false,
+    });
+
+    // Verify the symbol is attached correctly
+    expect((stepFn as any)[STEP_FUNCTION_NAME_SYMBOL]).toBe(stepName);
+  });
+
+  it('should not have STEP_FUNCTION_NAME_SYMBOL on regular functions', () => {
+    const regularFn = async (x: number) => x * 2;
+
+    // Regular functions should not have the symbol
+    expect((regularFn as any)[STEP_FUNCTION_NAME_SYMBOL]).toBeUndefined();
+  });
+
+  it('should lookup registered step function by name', () => {
+    const stepName = 'myRegisteredStep';
+    const stepFn = async (x: number) => x * 2;
+
+    // Register the step function
+    registerStepFunction(stepName, stepFn);
+
+    // Should be retrievable by name
+    const retrieved = getStepFunction(stepName);
+    expect(retrieved).toBe(stepFn);
+  });
+
+  it('should return undefined for non-existent registered step function', () => {
+    const retrieved = getStepFunction('nonExistentStep');
+    expect(retrieved).toBeUndefined();
+  });
+
+  it('should deserialize step function name through reviver', () => {
+    const stepName = 'testStep';
+    const stepFn = async () => 42;
+
+    // Register the step function
+    registerStepFunction(stepName, stepFn);
+
+    // Get the reviver and test it directly
+    const revivers = getCommonRevivers(vmGlobalThis);
+    const result = revivers.StepFunction(stepName);
+
+    expect(result).toBe(stepFn);
+  });
+
+  it('should throw error when reviver cannot find registered step function', () => {
+    const revivers = getCommonRevivers(vmGlobalThis);
+
+    let err: Error | undefined;
+    try {
+      revivers.StepFunction('nonExistentStep');
+    } catch (err_) {
+      err = err_ as Error;
+    }
+
+    expect(err).toBeDefined();
+    expect(err?.message).toContain('Step function "nonExistentStep" not found');
+    expect(err?.message).toContain('Make sure the step function is registered');
+  });
+});

packages/core/src/serialization.ts

@@ -1,8 +1,10 @@
 import { WorkflowRuntimeError } from '@workflow/errors';
 import * as devalue from 'devalue';
+import { getStepFunction } from './private.js';
 import { getWorld } from './runtime/world.js';
 import {
   BODY_INIT_SYMBOL,
+  STEP_FUNCTION_NAME_SYMBOL,
   STREAM_NAME_SYMBOL,
   STREAM_TYPE_SYMBOL,
   WEBHOOK_RESPONSE_WRITABLE,
@@ -169,6 +171,7 @@ export interface SerializableSpecial {
     redirected: boolean;
   };
   Set: any[];
+  StepFunction: string; // step function name/ID
   URL: string;
   URLSearchParams: string;
   Uint8Array: string; // base64 string
@@ -275,6 +278,12 @@ function getCommonReducers(global: Record<string, any> = globalThis) {
       };
     },
     Set: (value) => value instanceof global.Set && Array.from(value),
+    StepFunction: (value) => {
+      if (typeof value !== 'function') return false;
+      const stepName = value[STEP_FUNCTION_NAME_SYMBOL];
+      // Must return a truthy value; devalue will reject if we return false for a function
+      return stepName || false;
+    },
     URL: (value) => value instanceof global.URL && value.href,
     URLSearchParams: (value) => {
       if (!(value instanceof global.URLSearchParams)) return false;
@@ -464,7 +473,7 @@ function getStepReducers(
   };
 }
 
-function getCommonRevivers(global: Record<string, any> = globalThis) {
+export function getCommonRevivers(global: Record<string, any> = globalThis) {
   function reviveArrayBuffer(value: string) {
     // Handle sentinel value for zero-length buffers
     const base64 = value === '.' ? '' : value;
@@ -516,6 +525,15 @@ function getCommonRevivers(global: Record<string, any> = globalThis) {
     Map: (value) => new global.Map(value),
     RegExp: (value) => new global.RegExp(value.source, value.flags),
     Set: (value) => new global.Set(value),
+    StepFunction: (value) => {
+      const stepFn = getStepFunction(value);
+      if (!stepFn) {
+        throw new Error(
+          `Step function "${value}" not found. Make sure the step function is registered.`
+        );
+      }
+      return stepFn;
+    },
     URL: (value) => new global.URL(value),
     URLSearchParams: (value) =>
       new global.URLSearchParams(value === '.' ? '' : value),

packages/core/src/step.test.ts

@@ -5,6 +5,7 @@ import { monotonicFactory } from 'ulid';
 import { describe, expect, it, vi } from 'vitest';
 import { EventsConsumer } from './events-consumer.js';
 import { WorkflowSuspension } from './global.js';
+import { STEP_FUNCTION_NAME_SYMBOL } from './symbols.js';
 import type { WorkflowOrchestratorContext } from './private.js';
 import { createUseStep } from './step.js';
 import { createContext } from './vm/index.js';
@@ -175,4 +176,81 @@ describe('createUseStep', () => {
       ]
     `);
   });
+
+  it('should attach STEP_FUNCTION_NAME_SYMBOL to the returned function', async () => {
+    const ctx = setupWorkflowContext([
+      {
+        eventId: 'evnt_0',
+        runId: 'wrun_123',
+        eventType: 'step_completed',
+        correlationId: 'step_01K11TFZ62YS0YYFDQ3E8B9YCV',
+        eventData: {
+          result: [42],
+        },
+        createdAt: new Date(),
+      },
+    ]);
+
+    const useStep = createUseStep(ctx);
+    const myStep = useStep('calculateValue');
+
+    // The returned function should have the symbol attached with the step name
+    expect((myStep as any)[STEP_FUNCTION_NAME_SYMBOL]).toBe('calculateValue');
+  });
+
+  it('should allow the symbol to be non-writable and non-configurable', async () => {
+    const ctx = setupWorkflowContext([]);
+    const useStep = createUseStep(ctx);
+    const myStep = useStep('immutableStep');
+
+    const descriptor = Object.getOwnPropertyDescriptor(
+      myStep as any,
+      STEP_FUNCTION_NAME_SYMBOL
+    );
+
+    expect(descriptor).toBeDefined();
+    expect(descriptor?.value).toBe('immutableStep');
+    expect(descriptor?.writable).toBe(false);
+    expect(descriptor?.configurable).toBe(false);
+    expect(descriptor?.enumerable).toBe(false);
+  });
+
+  it('should attach symbol to multiple step references independently', async () => {
+    const ctx = setupWorkflowContext([
+      {
+        eventId: 'evnt_0',
+        runId: 'wrun_123',
+        eventType: 'step_completed',
+        correlationId: 'step_01K11TFZ62YS0YYFDQ3E8B9YCV',
+        eventData: {
+          result: [1],
+        },
+        createdAt: new Date(),
+      },
+      {
+        eventId: 'evnt_1',
+        runId: 'wrun_123',
+        eventType: 'step_completed',
+        correlationId: 'step_01K11TFZ62YS0YYFDQ3E8B9YCW',
+        eventData: {
+          result: [2],
+        },
+        createdAt: new Date(),
+      },
+    ]);
+
+    const useStep = createUseStep(ctx);
+    const stepA = useStep('stepA');
+    const stepB = useStep('stepB');
+
+    expect((stepA as any)[STEP_FUNCTION_NAME_SYMBOL]).toBe('stepA');
+    expect((stepB as any)[STEP_FUNCTION_NAME_SYMBOL]).toBe('stepB');
+
+    // Call them to verify they work independently
+    const resultA = stepA(1);
+    const resultB = stepB(2);
+
+    expect(await resultA).toBe(1);
+    expect(await resultB).toBe(2);
+  });
 });

packages/core/src/step.ts

@@ -6,12 +6,13 @@ import { stepLogger } from './logger.js';
 import type { WorkflowOrchestratorContext } from './private.js';
 import type { Serializable } from './schemas.js';
 import { hydrateStepReturnValue } from './serialization.js';
+import { STEP_FUNCTION_NAME_SYMBOL } from './symbols.js';
 
 export function createUseStep(ctx: WorkflowOrchestratorContext) {
   return function useStep<Args extends Serializable[], Result>(
     stepName: string
   ) {
-    return (...args: Args): Promise<Result> => {
+    const stepFunction = (...args: Args): Promise<Result> => {
       const { promise, resolve, reject } = withResolvers<Result>();
 
       const correlationId = `step_${ctx.generateUlid()}`;
@@ -124,5 +125,16 @@ export function createUseStep(ctx: WorkflowOrchestratorContext) {
 
       return promise;
     };
+
+    // Attach the step name to the function so it can be properly serialized
+    // when passed as an argument to another step
+    Object.defineProperty(stepFunction, STEP_FUNCTION_NAME_SYMBOL, {
+      value: stepName,
+      writable: false,
+      enumerable: false,
+      configurable: false,
+    });
+
+    return stepFunction;
   };
 }

packages/core/src/symbols.ts

@@ -9,3 +9,6 @@ export const BODY_INIT_SYMBOL = Symbol.for('BODY_INIT');
 export const WEBHOOK_RESPONSE_WRITABLE = Symbol.for(
   'WEBHOOK_RESPONSE_WRITABLE'
 );
+export const STEP_FUNCTION_NAME_SYMBOL = Symbol.for(
+  'WORKFLOW_STEP_FUNCTION_NAME'
+);

workbench/example/workflows/99_e2e.ts

@@ -466,3 +466,24 @@ export async function retryableAndFatalErrorWorkflow() {
 
   return { retryableResult, gotFatalError };
 }
+
+//////////////////////////////////////////////////////////
+
+async function stepWithStepFunctionArg(stepFn: (x: number) => Promise<number>) {
+  'use step';
+  // Call the passed step function reference
+  const result = await stepFn(10);
+  return result * 2;
+}
+
+async function doubleNumber(x: number) {
+  'use step';
+  return x * 2;
+}
+
+export async function stepFunctionPassingWorkflow() {
+  'use workflow';
+  // Pass a step function reference to another step
+  const result = await stepWithStepFunctionArg(doubleNumber);
+  return result;
+}