Support nested "use step" declarations in non-workflow functions

.changeset/sixty-carpets-join.md

@@ -0,0 +1,5 @@
+---
+"@workflow/swc-plugin": patch
+---
+
+Support nested "use step" declarations in non-workflow functions

packages/swc-plugin-workflow/transform/tests/errors/non-async-functions/output-client.js

@@ -9,13 +9,6 @@ export const syncWorkflow = ()=>{
     'use workflow';
     return 'test';
 };
-// Error: sync method with use step
-const obj = {
-    syncMethod () {
-        'use step';
-        return true;
-    }
-};
 // These are ok
 export async function validStep() {
     return 42;

packages/swc-plugin-workflow/transform/tests/errors/non-async-functions/output-workflow.js

@@ -9,13 +9,6 @@ export const syncWorkflow = ()=>{
     'use workflow';
     return 'test';
 };
-// Error: sync method with use step
-const obj = {
-    syncMethod () {
-        'use step';
-        return true;
-    }
-};
 // These are ok
 export var validStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//validStep");
 export const validWorkflow = async ()=>{

packages/swc-plugin-workflow/transform/tests/fixture/default-parameter-usage/input.js

@@ -0,0 +1,32 @@
+// Test case for functions used in default parameter values
+// The createDefaultDownloadFunction should NOT be removed by DCE
+
+const createDefaultDownloadFunction = (download = defaultDownload) => (requestedDownloads) =>
+  Promise.all(requestedDownloads.map(async (r) => r.isUrlSupportedByModel ? null : download(r)));
+
+async function defaultDownload(request) {
+  return fetch(request.url);
+}
+
+// This function uses createDefaultDownloadFunction in a default parameter value
+// DCE must NOT remove createDefaultDownloadFunction
+async function convertToLanguageModelPrompt({
+  prompt,
+  supportedUrls,
+  download = createDefaultDownloadFunction()
+}) {
+  return { prompt, supportedUrls, download };
+}
+
+export async function myWorkflow(input) {
+  'use workflow';
+
+  const result = await convertToLanguageModelPrompt({
+    prompt: input.prompt,
+    supportedUrls: {},
+    download: undefined
+  });
+
+  return result;
+}
+

packages/swc-plugin-workflow/transform/tests/fixture/default-parameter-usage/output-client.js

@@ -0,0 +1,7 @@
+// Test case for functions used in default parameter values
+// The createDefaultDownloadFunction should NOT be removed by DCE
+/**__internal_workflows{"workflows":{"input.js":{"myWorkflow":{"workflowId":"workflow//input.js//myWorkflow"}}}}*/;
+export async function myWorkflow(input) {
+    throw new Error("You attempted to execute workflow myWorkflow function directly. To start a workflow, use start(myWorkflow) from workflow/api");
+}
+myWorkflow.workflowId = "workflow//input.js//myWorkflow";

packages/swc-plugin-workflow/transform/tests/fixture/default-parameter-usage/output-step.js

@@ -0,0 +1,25 @@
+// Test case for functions used in default parameter values
+// The createDefaultDownloadFunction should NOT be removed by DCE
+/**__internal_workflows{"workflows":{"input.js":{"myWorkflow":{"workflowId":"workflow//input.js//myWorkflow"}}}}*/;
+const createDefaultDownloadFunction = (download = defaultDownload)=>(requestedDownloads)=>Promise.all(requestedDownloads.map(async (r)=>r.isUrlSupportedByModel ? null : download(r)));
+async function defaultDownload(request) {
+    return fetch(request.url);
+}
+// This function uses createDefaultDownloadFunction in a default parameter value
+// DCE must NOT remove createDefaultDownloadFunction
+async function convertToLanguageModelPrompt({ prompt, supportedUrls, download = createDefaultDownloadFunction() }) {
+    return {
+        prompt,
+        supportedUrls,
+        download
+    };
+}
+export async function myWorkflow(input) {
+    'use workflow';
+    const result = await convertToLanguageModelPrompt({
+        prompt: input.prompt,
+        supportedUrls: {},
+        download: undefined
+    });
+    return result;
+}

packages/swc-plugin-workflow/transform/tests/fixture/default-parameter-usage/output-workflow.js

@@ -0,0 +1,25 @@
+// Test case for functions used in default parameter values
+// The createDefaultDownloadFunction should NOT be removed by DCE
+/**__internal_workflows{"workflows":{"input.js":{"myWorkflow":{"workflowId":"workflow//input.js//myWorkflow"}}}}*/;
+const createDefaultDownloadFunction = (download = defaultDownload)=>(requestedDownloads)=>Promise.all(requestedDownloads.map(async (r)=>r.isUrlSupportedByModel ? null : download(r)));
+async function defaultDownload(request) {
+    return fetch(request.url);
+}
+// This function uses createDefaultDownloadFunction in a default parameter value
+// DCE must NOT remove createDefaultDownloadFunction
+async function convertToLanguageModelPrompt({ prompt, supportedUrls, download = createDefaultDownloadFunction() }) {
+    return {
+        prompt,
+        supportedUrls,
+        download
+    };
+}
+export async function myWorkflow(input) {
+    const result = await convertToLanguageModelPrompt({
+        prompt: input.prompt,
+        supportedUrls: {},
+        download: undefined
+    });
+    return result;
+}
+myWorkflow.workflowId = "workflow//input.js//myWorkflow";

packages/swc-plugin-workflow/transform/tests/fixture/module-level-step/output-client.js

@@ -1,7 +1,4 @@
 /**__internal_workflows{"steps":{"input.js":{"step":{"stepId":"step//input.js//step"},"stepArrow":{"stepId":"step//input.js//stepArrow"}}}}*/;
-const localArrow = async (input)=>{
-    return input.bar;
-};
 export async function step(input) {
     return input.foo;
 }

packages/swc-plugin-workflow/transform/tests/fixture/module-level-step/output-workflow.js

@@ -1,7 +1,4 @@
 /**__internal_workflows{"steps":{"input.js":{"step":{"stepId":"step//input.js//step"},"stepArrow":{"stepId":"step//input.js//stepArrow"}}}}*/;
 'use step';
-const localArrow = async (input)=>{
-    return input.bar;
-};
 export var step = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//step");
 export const stepArrow = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//stepArrow");

packages/swc-plugin-workflow/transform/tests/fixture/module-level-workflow/output-client.js

@@ -1,7 +1,4 @@
 /**__internal_workflows{"workflows":{"input.js":{"arrowWorkflow":{"workflowId":"workflow//input.js//arrowWorkflow"},"workflow":{"workflowId":"workflow//input.js//workflow"}}}}*/;
-const localArrow = async (input)=>{
-    return input.bar;
-};
 export async function workflow(input) {
     return input.foo;
 }

packages/swc-plugin-workflow/transform/tests/fixture/module-level-workflow/output-workflow.js

@@ -1,7 +1,4 @@
 /**__internal_workflows{"workflows":{"input.js":{"arrowWorkflow":{"workflowId":"workflow//input.js//arrowWorkflow"},"workflow":{"workflowId":"workflow//input.js//workflow"}}}}*/;
-const localArrow = async (input)=>{
-    return input.bar;
-};
 export async function workflow(input) {
     return input.foo;
 }

packages/swc-plugin-workflow/transform/tests/fixture/multiple-const-step-functions/input.js

@@ -14,3 +14,16 @@ export const fn3 = async () => {
   return 4;
 };
 
+// Test case: regular function BEFORE step function in same declaration
+// This verifies that processing doesn't skip the step function
+const regularArrow = () => 1, stepAfterRegular = async () => {
+  'use step';
+  return 5;
+};
+
+// Test case: regular function expression BEFORE step function
+const regularFn = function() { return 2; }, stepAfterRegularFn = async function() {
+  'use step';
+  return 6;
+};
+

packages/swc-plugin-workflow/transform/tests/fixture/multiple-const-step-functions/output-client.js

@@ -1,4 +1,4 @@
-/**__internal_workflows{"steps":{"input.js":{"fn1":{"stepId":"step//input.js//fn1"},"fn2":{"stepId":"step//input.js//fn2"},"fn3":{"stepId":"step//input.js//fn3"},"fn4":{"stepId":"step//input.js//fn4"}}}}*/;
+/**__internal_workflows{"steps":{"input.js":{"fn1":{"stepId":"step//input.js//fn1"},"fn2":{"stepId":"step//input.js//fn2"},"fn3":{"stepId":"step//input.js//fn3"},"fn4":{"stepId":"step//input.js//fn4"},"stepAfterRegular":{"stepId":"step//input.js//stepAfterRegular"},"stepAfterRegularFn":{"stepId":"step//input.js//stepAfterRegularFn"}}}}*/;
 const fn1 = async ()=>{
     return 1;
 }, fn2 = async ()=>{
@@ -9,3 +9,14 @@ export const fn3 = async ()=>{
 }, fn4 = async ()=>{
     return 4;
 };
+// Test case: regular function BEFORE step function in same declaration
+// This verifies that processing doesn't skip the step function
+const regularArrow = ()=>1, stepAfterRegular = async ()=>{
+    return 5;
+};
+// Test case: regular function expression BEFORE step function
+const regularFn = function() {
+    return 2;
+}, stepAfterRegularFn = async function() {
+    return 6;
+};

packages/swc-plugin-workflow/transform/tests/fixture/multiple-const-step-functions/output-step.js

@@ -1,5 +1,5 @@
 import { registerStepFunction } from "workflow/internal/private";
-/**__internal_workflows{"steps":{"input.js":{"fn1":{"stepId":"step//input.js//fn1"},"fn2":{"stepId":"step//input.js//fn2"},"fn3":{"stepId":"step//input.js//fn3"},"fn4":{"stepId":"step//input.js//fn4"}}}}*/;
+/**__internal_workflows{"steps":{"input.js":{"fn1":{"stepId":"step//input.js//fn1"},"fn2":{"stepId":"step//input.js//fn2"},"fn3":{"stepId":"step//input.js//fn3"},"fn4":{"stepId":"step//input.js//fn4"},"stepAfterRegular":{"stepId":"step//input.js//stepAfterRegular"},"stepAfterRegularFn":{"stepId":"step//input.js//stepAfterRegularFn"}}}}*/;
 const fn1 = async ()=>{
     return 1;
 }, fn2 = async ()=>{
@@ -10,7 +10,20 @@ export const fn3 = async ()=>{
 }, fn4 = async ()=>{
     return 4;
 };
+// Test case: regular function BEFORE step function in same declaration
+// This verifies that processing doesn't skip the step function
+const regularArrow = ()=>1, stepAfterRegular = async ()=>{
+    return 5;
+};
+// Test case: regular function expression BEFORE step function
+const regularFn = function() {
+    return 2;
+}, stepAfterRegularFn = async function() {
+    return 6;
+};
 registerStepFunction("step//input.js//fn1", fn1);
 registerStepFunction("step//input.js//fn2", fn2);
 registerStepFunction("step//input.js//fn3", fn3);
 registerStepFunction("step//input.js//fn4", fn4);
+registerStepFunction("step//input.js//stepAfterRegular", stepAfterRegular);
+registerStepFunction("step//input.js//stepAfterRegularFn", stepAfterRegularFn);

packages/swc-plugin-workflow/transform/tests/fixture/multiple-const-step-functions/output-workflow.js

@@ -1,3 +1,10 @@
-/**__internal_workflows{"steps":{"input.js":{"fn1":{"stepId":"step//input.js//fn1"},"fn2":{"stepId":"step//input.js//fn2"},"fn3":{"stepId":"step//input.js//fn3"},"fn4":{"stepId":"step//input.js//fn4"}}}}*/;
+/**__internal_workflows{"steps":{"input.js":{"fn1":{"stepId":"step//input.js//fn1"},"fn2":{"stepId":"step//input.js//fn2"},"fn3":{"stepId":"step//input.js//fn3"},"fn4":{"stepId":"step//input.js//fn4"},"stepAfterRegular":{"stepId":"step//input.js//stepAfterRegular"},"stepAfterRegularFn":{"stepId":"step//input.js//stepAfterRegularFn"}}}}*/;
 const fn1 = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//fn1"), fn2 = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//fn2");
 export const fn3 = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//fn3"), fn4 = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//fn4");
+// Test case: regular function BEFORE step function in same declaration
+// This verifies that processing doesn't skip the step function
+const regularArrow = ()=>1, stepAfterRegular = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//stepAfterRegular");
+// Test case: regular function expression BEFORE step function
+const regularFn = function() {
+    return 2;
+}, stepAfterRegularFn = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("step//input.js//stepAfterRegularFn");

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-in-workflow/output-step.js

@@ -11,6 +11,14 @@ var helpers$objectStep = async (x, y)=>{
     return x + y + 10;
 };
 export async function example(a, b) {
+    "use workflow";
+    const step = example$step;
+    // Arrow function with const
+    const arrowStep = example$arrowStep;
+    // Arrow function with let
+    let letArrowStep = example$letArrowStep;
+    // Arrow function with var
+    var varArrowStep = example$varArrowStep;
     // Object with step method
     const helpers = {
         objectStep: helpers$objectStep

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/input.js

@@ -1,6 +1,67 @@
 import { DurableAgent } from '@workflow/ai/agent';
 import { gateway } from 'ai';
 
+function stepWrapperReturnArrowFunctionVar(a, b, c) {
+  const fn = async () => {
+    'use step';
+    return a + b + c;
+  };
+  return fn;
+}
+
+function stepWrapperReturnNamedFunction(a, b, c) {
+  return async function f() {
+    'use step';
+    return a + b + c;
+  };
+}
+
+function stepWrapperReturnArrowFunction(a, b, c) {
+  return async () => {
+    'use step';
+    return a + b + c;
+  };
+}
+
+function stepWrapperReturnNamedFunctionVar(a, b, c) {
+  async function fn() {
+    'use step';
+    return a + b + c;
+  }
+  return fn;
+}
+
+const arrowWrapperReturnArrowFunctionVar = (a, b, c) => {
+  const fn = async () => {
+    'use step';
+    return a + b + c;
+  };
+  return fn;
+}
+
+const arrowWrapperReturnNamedFunction = (a, b, c) => {
+  return async function f() {
+    'use step';
+    return a + b + c;
+  };
+}
+
+const arrowWrapperReturnArrowFunction = (a, b, c) => {
+  return async () => {
+    'use step';
+    return a + b + c;
+  };
+}
+
+const arrowWrapperReturnNamedFunctionVar = (a, b, c) => {
+  async function fn() {
+    'use step';
+    return a + b + c;
+  }
+  return fn;
+}
+
+
 export async function wflow() {
   'use workflow';
   let count = 42;
@@ -27,4 +88,13 @@ export async function wflow() {
       console.log('count', count);
     },
   });
+
+  await stepWrapperReturnArrowFunctionVar(1, 2, 3)();
+  await stepWrapperReturnNamedFunction(1, 2, 3)();
+  await stepWrapperReturnArrowFunction(1, 2, 3)();
+  await stepWrapperReturnNamedFunctionVar(1, 2, 3)();
+  await arrowWrapperReturnArrowFunctionVar(1, 2, 3)();
+  await arrowWrapperReturnNamedFunction(1, 2, 3)();
+  await arrowWrapperReturnArrowFunction(1, 2, 3)();
+  await arrowWrapperReturnNamedFunctionVar(1, 2, 3)();
 }

packages/swc-plugin-workflow/transform/tests/fixture/nested-step-with-closure/output-client.js

@@ -1,4 +1,4 @@
-/**__internal_workflows{"workflows":{"input.js":{"wflow":{"workflowId":"workflow//input.js//wflow"}}}}*/;
+/**__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"},"f":{"stepId":"step//input.js//f"},"fn":{"stepId":"step//input.js//fn"}}}}*/;
 export async function wflow() {
     throw new Error("You attempted to execute workflow wflow function directly. To start a workflow, use start(wflow) from workflow/api");
 }