diff --git a/.changeset/bright-ducks-travel.md b/.changeset/bright-ducks-travel.md new file mode 100644 index 000000000..2da0a13bc --- /dev/null +++ b/.changeset/bright-ducks-travel.md @@ -0,0 +1,5 @@ +--- +"@workflow/core": patch +--- + +Support closure variables for serialized step functions diff --git a/packages/core/e2e/e2e.test.ts b/packages/core/e2e/e2e.test.ts index 72991d163..92336b08e 100644 --- a/packages/core/e2e/e2e.test.ts +++ b/packages/core/e2e/e2e.test.ts @@ -707,7 +707,7 @@ describe('e2e', () => { ); 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 @@ describe('e2e', () => { } ); + 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 }, diff --git a/packages/core/src/serialization.test.ts b/packages/core/src/serialization.test.ts index 36f11b479..2b38a655c 100644 --- a/packages/core/src/serialization.test.ts +++ b/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 }); }); }); diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index 57fa8d9ae..af81af68a 100644 --- a/packages/core/src/serialization.ts +++ b/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; + }; URL: string; URLSearchParams: string; Uint8Array: string; // base64 string @@ -291,7 +295,18 @@ function getCommonReducers(global: Record = 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 = 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), diff --git a/packages/core/src/step.ts b/packages/core/src/step.ts index bf0fc1cc0..f5d0240cb 100644 --- a/packages/core/src/step.ts +++ b/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; }; } diff --git a/packages/swc-plugin-workflow/transform/src/lib.rs b/packages/swc-plugin-workflow/transform/src/lib.rs index 134c02289..c48e675e1 100644 --- a/packages/swc-plugin-workflow/transform/src/lib.rs +++ b/packages/swc-plugin-workflow/transform/src/lib.rs @@ -170,6 +170,8 @@ pub struct StepTransform { // Track context for validation in_step_function: bool, in_workflow_function: bool, + // Track the current workflow function name (for nested step naming) + current_workflow_function_name: Option, // Track workflow functions that need to be expanded into multiple exports workflow_exports_to_expand: Vec<(String, Expr, swc_core::common::Span)>, // Track workflow functions that need workflowId property in client mode @@ -182,8 +184,8 @@ pub struct StepTransform { // (parent_var_name, prop_name, arrow_expr, span) object_property_step_functions: Vec<(String, String, ArrowExpr, swc_core::common::Span)>, // Track nested step functions inside workflow functions for hoisting in step mode - // (fn_name, fn_expr, span, closure_vars, was_arrow) - nested_step_functions: Vec<(String, FnExpr, swc_core::common::Span, Vec, bool)>, + // (fn_name, fn_expr, span, closure_vars, was_arrow, parent_workflow_name) + nested_step_functions: Vec<(String, FnExpr, swc_core::common::Span, Vec, bool, String)>, // Counter for anonymous function names #[allow(dead_code)] anonymous_fn_counter: usize, @@ -714,13 +716,20 @@ impl StepTransform { fn_decl.function.span, closure_vars, false, // Regular function, not arrow + self.current_workflow_function_name.clone().unwrap_or_default(), )); *stmt = Stmt::Empty(EmptyStmt { span: DUMMY_SP }); return; } TransformMode::Workflow => { + // Include parent workflow name in step ID + let step_fn_name = if let Some(parent) = &self.current_workflow_function_name { + format!("{}/{}", parent, fn_name) + } else { + fn_name.clone() + }; let step_id = self.create_id( - Some(&fn_name), + Some(&step_fn_name), fn_decl.function.span, false, ); @@ -915,6 +924,7 @@ impl StepTransform { in_callee: false, in_step_function: false, in_workflow_function: false, + current_workflow_function_name: None, workflow_exports_to_expand: Vec::new(), workflow_functions_needing_id: Vec::new(), step_exports_to_convert: Vec::new(), @@ -2639,7 +2649,7 @@ impl VisitMut for StepTransform { let needs_closure_import = self .nested_step_functions .iter() - .any(|(_, _, _, closure_vars, _)| !closure_vars.is_empty()); + .any(|(_, _, _, closure_vars, _, _)| !closure_vars.is_empty()); if needs_register_import || needs_closure_import { imports_to_add.push(self.create_private_imports( @@ -2672,8 +2682,14 @@ impl VisitMut for StepTransform { // Process nested step functions FIRST (they typically appear earlier in source) let nested_functions: Vec<_> = self.nested_step_functions.drain(..).collect(); - - for (fn_name, mut fn_expr, span, closure_vars, was_arrow) in nested_functions { + + for (fn_name, mut fn_expr, span, closure_vars, was_arrow, parent_workflow_name) in nested_functions { + // Generate hoisted name including parent workflow function name + let hoisted_name = if parent_workflow_name.is_empty() { + fn_name.clone() + } else { + format!("{}${}", parent_workflow_name, fn_name) + }; // If there are closure variables, add destructuring as first statement if !closure_vars.is_empty() { if let Some(body) = &mut fn_expr.function.body { @@ -2742,7 +2758,7 @@ impl VisitMut for StepTransform { span: DUMMY_SP, name: Pat::Ident(BindingIdent { id: Ident::new( - fn_name.clone().into(), + hoisted_name.clone().into(), DUMMY_SP, SyntaxContext::empty(), ), @@ -2757,7 +2773,7 @@ impl VisitMut for StepTransform { // Keep as function declaration: async function name() { ... } ModuleItem::Stmt(Stmt::Decl(Decl::Fn(FnDecl { ident: Ident::new( - fn_name.clone().into(), + hoisted_name.clone().into(), DUMMY_SP, SyntaxContext::empty(), ), @@ -2770,8 +2786,13 @@ impl VisitMut for StepTransform { module.body.insert(current_insert_pos, hoisted_decl); current_insert_pos += 1; - // Create a registration call - let step_id = self.create_id(Some(&fn_name), span, false); + // Create a registration call with parent workflow name in the step ID + let step_fn_name = if parent_workflow_name.is_empty() { + fn_name.clone() + } else { + format!("{}/{}", parent_workflow_name, fn_name) + }; + let step_id = self.create_id(Some(&step_fn_name), span, false); let registration_call = Stmt::Expr(ExprStmt { span: DUMMY_SP, expr: Box::new(Expr::Call(CallExpr { @@ -2794,7 +2815,7 @@ impl VisitMut for StepTransform { ExprOrSpread { spread: None, expr: Box::new(Expr::Ident(Ident::new( - fn_name.into(), + hoisted_name.into(), DUMMY_SP, SyntaxContext::empty(), ))), @@ -3000,6 +3021,7 @@ impl VisitMut for StepTransform { // Set context for forbidden expression checking let old_in_step = self.in_step_function; let old_in_workflow = self.in_workflow_function; + let old_workflow_name = self.current_workflow_function_name.clone(); let old_in_module = self.in_module_level; if has_step_directive { @@ -3016,6 +3038,7 @@ impl VisitMut for StepTransform { // Restore context self.in_step_function = old_in_step; self.in_workflow_function = old_in_workflow; + self.current_workflow_function_name = old_workflow_name; self.in_module_level = old_in_module; } @@ -3026,6 +3049,7 @@ impl VisitMut for StepTransform { // Set context for forbidden expression checking let old_in_step = self.in_step_function; let old_in_workflow = self.in_workflow_function; + let old_workflow_name = self.current_workflow_function_name.clone(); let old_in_module = self.in_module_level; if has_step_directive { @@ -3042,6 +3066,7 @@ impl VisitMut for StepTransform { // Restore context self.in_step_function = old_in_step; self.in_workflow_function = old_in_workflow; + self.current_workflow_function_name = old_workflow_name; self.in_module_level = old_in_module; } @@ -3747,8 +3772,13 @@ impl VisitMut for StepTransform { } let old_in_workflow = self.in_workflow_function; + let old_workflow_name = self.current_workflow_function_name.clone(); if is_workflow_function { self.in_workflow_function = true; + // Get the function name for context tracking + if let Decl::Fn(fn_decl) = &export_decl.decl { + self.current_workflow_function_name = Some(fn_decl.ident.sym.to_string()); + } } match &mut export_decl.decl { @@ -4234,6 +4264,7 @@ impl VisitMut for StepTransform { // Restore in_workflow_function flag self.in_workflow_function = old_in_workflow; + self.current_workflow_function_name = old_workflow_name; } fn visit_mut_var_decl(&mut self, var_decl: &mut VarDecl) { @@ -4460,6 +4491,7 @@ impl VisitMut for StepTransform { arrow_expr.span, closure_vars, true, // Was an arrow function + self.current_workflow_function_name.clone().unwrap_or_default(), )); // Mark the entire var declarator for removal by nulling out the init @@ -4469,8 +4501,14 @@ impl VisitMut for StepTransform { } TransformMode::Workflow => { // Replace with proxy reference (not a function call) + // Include parent workflow name in step ID + let step_fn_name = if let Some(parent) = &self.current_workflow_function_name { + format!("{}/{}", parent, name) + } else { + name.clone() + }; let step_id = self.create_id( - Some(&name), + Some(&step_fn_name), arrow_expr.span, false, ); @@ -4995,6 +5033,7 @@ impl VisitMut for StepTransform { arrow_expr.span, closure_vars, true, // Was an arrow function + self.current_workflow_function_name.clone().unwrap_or_default(), )); // Replace with identifier reference @@ -5009,8 +5048,14 @@ impl VisitMut for StepTransform { self.remove_use_step_directive_arrow( &mut arrow_expr.body, ); + // Include parent workflow name in step ID + let step_fn_name = if let Some(parent) = &self.current_workflow_function_name { + format!("{}/{}", parent, generated_name) + } else { + generated_name.clone() + }; let step_id = self.create_id( - Some(&generated_name), + Some(&step_fn_name), arrow_expr.span, false, ); @@ -5076,6 +5121,7 @@ impl VisitMut for StepTransform { fn_expr.function.span, closure_vars, false, // Was a function expression + self.current_workflow_function_name.clone().unwrap_or_default(), )); // Replace with identifier reference @@ -5090,8 +5136,14 @@ impl VisitMut for StepTransform { self.remove_use_step_directive( &mut fn_expr.function.body, ); + // Include parent workflow name in step ID + let step_fn_name = if let Some(parent) = &self.current_workflow_function_name { + format!("{}/{}", parent, generated_name) + } else { + generated_name.clone() + }; let step_id = self.create_id( - Some(&generated_name), + Some(&step_fn_name), fn_expr.function.span, false, ); @@ -5164,6 +5216,7 @@ impl VisitMut for StepTransform { method_prop.function.span, closure_vars, false, // Was a method + self.current_workflow_function_name.clone().unwrap_or_default(), )); // Replace method with property pointing to identifier @@ -5182,8 +5235,14 @@ impl VisitMut for StepTransform { self.remove_use_step_directive( &mut method_prop.function.body, ); + // Include parent workflow name in step ID + let step_fn_name = if let Some(parent) = &self.current_workflow_function_name { + format!("{}/{}", parent, generated_name) + } else { + generated_name.clone() + }; let step_id = self.create_id( - Some(&generated_name), + Some(&step_fn_name), method_prop.function.span, false, ); diff --git a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-step.js b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-step.js index 9022b31a7..594a8196a 100644 --- a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-step.js +++ b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-step.js @@ -1,12 +1,12 @@ import { registerStepFunction } from "workflow/internal/private"; /**__internal_workflows{"workflows":{"input.js":{"example":{"workflowId":"workflow//input.js//example"}}},"steps":{"input.js":{"arrowStep":{"stepId":"step//input.js//arrowStep"},"helpers/objectStep":{"stepId":"step//input.js//helpers/objectStep"},"letArrowStep":{"stepId":"step//input.js//letArrowStep"},"step":{"stepId":"step//input.js//step"},"varArrowStep":{"stepId":"step//input.js//varArrowStep"}}}}*/; // Function declaration step -async function step(a, b) { +async function example$step(a, b) { return a + b; } -var arrowStep = async (x, y)=>x * y; -var letArrowStep = async (x, y)=>x - y; -var varArrowStep = async (x, y)=>x / y; +var example$arrowStep = async (x, y)=>x * y; +var example$letArrowStep = async (x, y)=>x - y; +var example$varArrowStep = async (x, y)=>x / y; var helpers$objectStep = async (x, y)=>{ return x + y + 10; }; @@ -22,8 +22,8 @@ export async function example(a, b) { const val5 = await helpers.objectStep(a, b); return val + val2 + val3 + val4 + val5; } -registerStepFunction("step//input.js//step", step); -registerStepFunction("step//input.js//arrowStep", arrowStep); -registerStepFunction("step//input.js//letArrowStep", letArrowStep); -registerStepFunction("step//input.js//varArrowStep", varArrowStep); +registerStepFunction("step//input.js//example/step", example$step); +registerStepFunction("step//input.js//example/arrowStep", example$arrowStep); +registerStepFunction("step//input.js//example/letArrowStep", example$letArrowStep); +registerStepFunction("step//input.js//example/varArrowStep", example$varArrowStep); registerStepFunction("step//input.js//helpers/objectStep", helpers$objectStep); diff --git a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-workflow.js b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-workflow.js index f530b39b7..f6ed62968 100644 --- a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-workflow.js +++ b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-workflow.js @@ -1,12 +1,12 @@ /**__internal_workflows{"workflows":{"input.js":{"example":{"workflowId":"workflow//input.js//example"}}},"steps":{"input.js":{"arrowStep":{"stepId":"step//input.js//arrowStep"},"helpers/objectStep":{"stepId":"step//input.js//helpers/objectStep"},"letArrowStep":{"stepId":"step//input.js//letArrowStep"},"step":{"stepId":"step//input.js//step"},"varArrowStep":{"stepId":"step//input.js//varArrowStep"}}}}*/; export async function example(a, b) { - var step = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//step"); + var step = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//example/step"); // Arrow function with const - const arrowStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//arrowStep"); + const arrowStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//example/arrowStep"); // Arrow function with let - let letArrowStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//letArrowStep"); + let letArrowStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//example/letArrowStep"); // Arrow function with var - var varArrowStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//varArrowStep"); + var varArrowStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//example/varArrowStep"); // Object with step method const helpers = { objectStep: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//helpers/objectStep") diff --git a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-step.js b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-step.js index 82c9b1509..531693db2 100644 --- a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-step.js +++ b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-step.js @@ -2,20 +2,20 @@ import { __private_getClosureVars, registerStepFunction } from "workflow/interna import { DurableAgent } from '@workflow/ai/agent'; import { gateway } from 'ai'; /**__internal_workflows{"workflows":{"input.js":{"wflow":{"workflowId":"workflow//input.js//wflow"}}},"steps":{"input.js":{"_anonymousStep0":{"stepId":"step//input.js//_anonymousStep0"},"_anonymousStep1":{"stepId":"step//input.js//_anonymousStep1"},"_anonymousStep2":{"stepId":"step//input.js//_anonymousStep2"},"namedStepWithClosureVars":{"stepId":"step//input.js//namedStepWithClosureVars"}}}}*/; -async function namedStepWithClosureVars() { +async function wflow$namedStepWithClosureVars() { const { count } = __private_getClosureVars(); console.log('count', count); } -var _anonymousStep0 = async ()=>{ +var wflow$_anonymousStep0 = async ()=>{ const { count } = __private_getClosureVars(); console.log('count', count); return gateway('openai/gpt-5'); }; -async function _anonymousStep1() { +async function wflow$_anonymousStep1() { const { count } = __private_getClosureVars(); console.log('count', count); } -async function _anonymousStep2() { +async function wflow$_anonymousStep2() { const { count } = __private_getClosureVars(); console.log('count', count); } @@ -27,7 +27,7 @@ export async function wflow() { methodWithClosureVars: _anonymousStep2 }); } -registerStepFunction("step//input.js//namedStepWithClosureVars", namedStepWithClosureVars); -registerStepFunction("step//input.js//_anonymousStep0", _anonymousStep0); -registerStepFunction("step//input.js//_anonymousStep1", _anonymousStep1); -registerStepFunction("step//input.js//_anonymousStep2", _anonymousStep2); +registerStepFunction("step//input.js//wflow/namedStepWithClosureVars", wflow$namedStepWithClosureVars); +registerStepFunction("step//input.js//wflow/_anonymousStep0", wflow$_anonymousStep0); +registerStepFunction("step//input.js//wflow/_anonymousStep1", wflow$_anonymousStep1); +registerStepFunction("step//input.js//wflow/_anonymousStep2", wflow$_anonymousStep2); diff --git a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-workflow.js b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-workflow.js index f9e6391de..6e768914f 100644 --- a/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-workflow.js +++ b/packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-workflow.js @@ -2,17 +2,17 @@ import { DurableAgent } from '@workflow/ai/agent'; /**__internal_workflows{"workflows":{"input.js":{"wflow":{"workflowId":"workflow//input.js//wflow"}}},"steps":{"input.js":{"_anonymousStep0":{"stepId":"step//input.js//_anonymousStep0"},"_anonymousStep1":{"stepId":"step//input.js//_anonymousStep1"},"_anonymousStep2":{"stepId":"step//input.js//_anonymousStep2"},"namedStepWithClosureVars":{"stepId":"step//input.js//namedStepWithClosureVars"}}}}*/; export async function wflow() { let count = 42; - var namedStepWithClosureVars = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//namedStepWithClosureVars", ()=>({ + var namedStepWithClosureVars = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//wflow/namedStepWithClosureVars", ()=>({ count })); const agent = new DurableAgent({ - arrowFunctionWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep0", ()=>({ + arrowFunctionWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//wflow/_anonymousStep0", ()=>({ count })), - namedFunctionWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep1", ()=>({ + namedFunctionWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//wflow/_anonymousStep1", ()=>({ count })), - methodWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep2", ()=>({ + methodWithClosureVars: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//wflow/_anonymousStep2", ()=>({ count })) }); diff --git a/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-step.js b/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-step.js index b1cd46531..a6b2e456e 100644 --- a/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-step.js +++ b/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-step.js @@ -3,8 +3,8 @@ import { DurableAgent } from '@workflow/ai/agent'; import { gateway, tool } from 'ai'; import * as z from 'zod'; /**__internal_workflows{"workflows":{"input.js":{"test":{"workflowId":"workflow//input.js//test"}}},"steps":{"input.js":{"_anonymousStep0":{"stepId":"step//input.js//_anonymousStep0"},"_anonymousStep1":{"stepId":"step//input.js//_anonymousStep1"}}}}*/; -var _anonymousStep0 = async ()=>gateway('openai/gpt-5'); -var _anonymousStep1 = async ({ location })=>`Weather in ${location}: Sunny, 72°F`; +var test$_anonymousStep0 = async ()=>gateway('openai/gpt-5'); +var test$_anonymousStep1 = async ({ location })=>`Weather in ${location}: Sunny, 72°F`; export async function test() { 'use workflow'; const agent = new DurableAgent({ @@ -28,5 +28,5 @@ export async function test() { ] }); } -registerStepFunction("step//input.js//_anonymousStep0", _anonymousStep0); -registerStepFunction("step//input.js//_anonymousStep1", _anonymousStep1); +registerStepFunction("step//input.js//test/_anonymousStep0", test$_anonymousStep0); +registerStepFunction("step//input.js//test/_anonymousStep1", test$_anonymousStep1); diff --git a/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-workflow.js b/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-workflow.js index 5836e426c..915a60f89 100644 --- a/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-workflow.js +++ b/packages/swc-plugin-workflow/transform/tests/fixture/nested-steps-in-object-constructor/output-workflow.js @@ -4,14 +4,14 @@ import * as z from 'zod'; /**__internal_workflows{"workflows":{"input.js":{"test":{"workflowId":"workflow//input.js//test"}}},"steps":{"input.js":{"_anonymousStep0":{"stepId":"step//input.js//_anonymousStep0"},"_anonymousStep1":{"stepId":"step//input.js//_anonymousStep1"}}}}*/; export async function test() { const agent = new DurableAgent({ - model: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep0"), + model: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//test/_anonymousStep0"), tools: { getWeather: tool({ description: 'Get weather for a location', inputSchema: z.object({ location: z.string() }), - execute: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//_anonymousStep1") + execute: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//test/_anonymousStep1") }) } }); diff --git a/workbench/example/workflows/99_e2e.ts b/workbench/example/workflows/99_e2e.ts index 7e442e52e..df5e0b22c 100644 --- a/workbench/example/workflows/99_e2e.ts +++ b/workbench/example/workflows/99_e2e.ts @@ -496,7 +496,7 @@ export async function hookCleanupTestWorkflow( export async function stepFunctionPassingWorkflow() { 'use workflow'; - // Pass a step function reference to another step + // Pass a step function reference to another step (without closure vars) const result = await stepWithStepFunctionArg(doubleNumber); return result; } @@ -515,8 +515,37 @@ async function doubleNumber(x: number) { ////////////////////////////////////////////////////////// +export async function stepFunctionWithClosureWorkflow() { + 'use workflow'; + const multiplier = 3; + const prefix = 'Result: '; + + // Create a step function that captures closure variables + const calculate = async (x: number) => { + 'use step'; + return `${prefix}${x * multiplier}`; + }; + + // Pass the step function (with closure vars) to another step + const result = await stepThatCallsStepFn(calculate, 7); + return result; +} + +async function stepThatCallsStepFn( + stepFn: (x: number) => Promise, + value: number +) { + 'use step'; + // Call the passed step function - closure vars should be preserved + const result = await stepFn(value); + return `Wrapped: ${result}`; +} + +////////////////////////////////////////////////////////// + export async function closureVariableWorkflow(baseValue: number) { 'use workflow'; + // biome-ignore lint/style/useConst: Intentionally using `let` instead of `const` let multiplier = 3; const prefix = 'Result: ';