Apply SWC transformation on step functions returned from factory function

Author: TooTallNate

.changeset/metal-cycles-slide.md

@@ -0,0 +1,5 @@
+---
+"@workflow/swc-plugin": patch
+---
+
+Apply SWC transformation on step functions returned from factory function

packages/swc-plugin-workflow/transform/src/lib.rs

@@ -786,7 +786,40 @@ impl StepTransform {
                     stmt.visit_mut_children_with(self);
                 }
             }
-            Stmt::Decl(Decl::Var(_)) => {
+            Stmt::Decl(Decl::Var(var_decl)) => {
+                // Check if any declarators contain arrow functions with object literal bodies
+                for declarator in &mut var_decl.decls {
+                    if let Some(init) = &mut declarator.init {
+                        if let Pat::Ident(binding) = &declarator.name {
+                            let name = binding.id.sym.to_string();
+                            
+                            // Check if the initializer is an arrow function with object literal body
+                            if let Expr::Arrow(arrow_expr) = &mut **init {
+                                match &mut *arrow_expr.body {
+                                    BlockStmtOrExpr::Expr(expr) => {
+                                        // Handle both direct object literals and parenthesized ones
+                                        let obj_lit_mut = match &mut **expr {
+                                            Expr::Object(obj) => Some(obj),
+                                            Expr::Paren(paren) => {
+                                                if let Expr::Object(obj) = &mut *paren.expr {
+                                                    Some(obj)
+                                                } else {
+                                                    None
+                                                }
+                                            }
+                                            _ => None,
+                                        };
+                                        
+                                        if let Some(obj_lit) = obj_lit_mut {
+                                            self.process_object_properties_for_step_functions(obj_lit, &name);
+                                        }
+                                    }
+                                    _ => {}
+                                }
+                            }
+                        }
+                    }
+                }
                 stmt.visit_mut_children_with(self);
             }
             _ => {

packages/swc-plugin-workflow/transform/tests/fixture/factory-with-step-method/output-client.js

@@ -1,7 +1,6 @@
 import fs from 'fs/promises';
 const myFactory = ()=>({
         myStep: async ()=>{
-            'use step';
             await fs.mkdir('test');
         }
     });

packages/swc-plugin-workflow/transform/tests/fixture/factory-with-step-method/output-step.js

@@ -1,8 +1,11 @@
+import { registerStepFunction } from "workflow/internal/private";
 import fs from 'fs/promises';
+/**__internal_workflows{"steps":{"input.js":{"myFactory/myStep":{"stepId":"step//input.js//myFactory/myStep"}}}}*/;
+var myFactory$myStep = async ()=>{
+    await fs.mkdir('test');
+};
 const myFactory = ()=>({
-        myStep: async ()=>{
-            'use step';
-            await fs.mkdir('test');
-        }
+        myStep: myFactory$myStep
     });
 export default myFactory;
+registerStepFunction("step//input.js//myFactory/myStep", myFactory$myStep);

packages/swc-plugin-workflow/transform/tests/fixture/factory-with-step-method/output-workflow.js

@@ -1,8 +1,5 @@
-import fs from 'fs/promises';
+/**__internal_workflows{"steps":{"input.js":{"myFactory/myStep":{"stepId":"step//input.js//myFactory/myStep"}}}}*/;
 const myFactory = ()=>({
-        myStep: async ()=>{
-            'use step';
-            await fs.mkdir('test');
-        }
+        myStep: globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//myFactory/myStep")
     });
 export default myFactory;