Support passing step function references across serialization boundaries

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.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,11 @@ 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];
+      return stepName ? stepName : false;
+    },
     URL: (value) => value instanceof global.URL && value.href,
     URLSearchParams: (value) => {
       if (!(value instanceof global.URLSearchParams)) return false;
@@ -516,6 +524,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.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'
+);