@output.ai/core 0.1.9-dev.pr156.0 → 0.1.9

package/src/worker/webpack_loaders/tools.js

@@ -6,11 +6,14 @@ import {
   callExpression,
   functionExpression,
   identifier,
+  isArrowFunctionExpression,
   isAssignmentPattern,
   isBlockStatement,
   isCallExpression,
   isExportNamedDeclaration,
+  isFunctionExpression,
   isIdentifier,
+  isVariableDeclarator,
   isStringLiteral,
   isVariableDeclaration,
   isObjectExpression,
@@ -18,7 +21,8 @@ import {
   returnStatement,
   stringLiteral,
   thisExpression,
-  isExportDefaultDeclaration
+  isExportDefaultDeclaration,
+  isFunctionDeclaration
 } from '@babel/types';
 import { ComponentFile, EXTRANEOUS_FILE, ExtraneousFileList, NodeType } from './consts.js';
 
@@ -170,6 +174,28 @@ export const getFileKind = path => {
 export const createThisMethodCall = ( method, literalName, args ) =>
   callExpression( memberExpression( thisExpression(), identifier( method ) ), [ stringLiteral( literalName ), ...args ] );
 
+/**
+ * Build a CallExpression that binds `this` at the call site:
+ *   fn(arg1, arg2) -> fn.call(this, arg1, arg2)
+ *
+ * When to use:
+ * - Inside workflow `fn` rewriting, local call-chain functions must receive the dynamic `this`
+ *   so that emitted `this.invokeStep(...)` and similar calls inside them operate correctly.
+ *
+ * Example:
+ *   // Input AST intent:
+ *   foo(a, b);
+ *
+ *   // Rewritten AST:
+ *   foo.call(this, a, b);
+ *
+ * @param {string} calleeName - Identifier name of the function being called (e.g., 'foo').
+ * @param {import('@babel/types').Expression[]} args - Original call arguments.
+ * @returns {import('@babel/types').CallExpression} CallExpression node representing `callee.call(this, ...args)`.
+ */
+export const bindThisAtCallSite = ( calleeName, args ) =>
+  callExpression( memberExpression( identifier( calleeName ), identifier( 'call' ) ), [ thisExpression(), ...args ] );
+
 /**
  * Resolve an options object's name property to a string.
  * Accepts literal strings or top-level const string identifiers.
@@ -332,3 +358,93 @@ export const buildWorkflowNameMap = ( path, cache ) => {
   cache.set( path, result );
   return result;
 };
+
+/**
+ * Determine whether a node represents a function body usable as a workflow `fn`.
+ *
+ * Why this matters:
+ * - Workflow `fn` needs a dynamic `this` so the rewriter can emit calls like `this.invokeStep(...)`.
+ * - Arrow functions do not have their own `this`; they capture `this` lexically, which breaks the runtime contract.
+ *
+ * Accepts:
+ * - FunctionExpression (possibly async/generator), e.g.:
+ *   const obj = {
+ *     fn: async function (input) {
+ *       return input;
+ *     }
+ *   };
+ *
+ * Rejects:
+ * - ArrowFunctionExpression, e.g.:
+ *   const obj = {
+ *     fn: async (input) => input
+ *   };
+ *
+ * - Any other non-function expression.
+ *
+ * Notes:
+ * - The rewriter will proactively convert arrow `fn` to a FunctionExpression before further processing.
+ *
+ * @param {import('@babel/types').Expression} v - Candidate node for `fn` value.
+ * @returns {boolean} True if `v` is a FunctionExpression and not an arrow function.
+ */
+export const isFunction = v => isFunctionExpression( v ) && !isArrowFunctionExpression( v );
+
+/**
+ * Determine whether a variable declarator represents a function-like value.
+ *
+ * Use case:
+ * - When `fn` calls a locally-declared function (directly or transitively), we need to:
+ *   - propagate `this` to that function call (`callee.call(this, ...)`)
+ *   - traverse into that function's body to rewrite imported step/workflow/evaluator calls.
+ *
+ * Matches patterns like:
+ * - Function expression:
+ *   const foo = function (x) { return x + 1; };
+ *
+ * - Async/generator function expression:
+ *   const foo = async function (x) { return await work(x); };
+ *
+ * - Arrow function (will be normalized to FunctionExpression by the rewriter):
+ *   const foo = (x) => x + 1;
+ *   const foo = async (x) => await work(x);
+ *
+ * Does not match:
+ * - Non-function initializers:
+ *   const foo = 42;
+ *   const foo = someIdentifier;
+ *
+ * @param {import('@babel/types').Node} v - AST node (typically a VariableDeclarator).
+ * @returns {boolean} True if the declarator's initializer is a function (arrow or function expression).
+ */
+export const isVarFunction = v =>
+  isVariableDeclarator( v ) && ( isFunctionExpression( v.init ) || isArrowFunctionExpression( v.init ) );
+
+/**
+ * Determine whether a binding node corresponds to a function-like declaration usable
+ * as a call-chain function target during workflow rewriting.
+ *
+ * Matches:
+ * - FunctionDeclaration:
+ *     function foo(x) { return x + 1; }
+ *
+ * - VariableDeclarator initialized with a function or arrow (normalized later):
+ *     const foo = function (x) { return x + 1; };
+ *     const foo = (x) => x + 1;
+ *
+ * Non-matches:
+ * - Any binding that is not a function declaration nor a variable declarator with a function initializer.
+ *
+ * Why this matters:
+ * - The rewriter traverses call chains from the workflow `fn`. It must recognize which local
+ *   callees are valid function bodies to rewrite and into which it can propagate `this`.
+ *
+ * @param {import('@babel/types').Node} node - Binding path node (FunctionDeclaration or VariableDeclarator).
+ * @returns {boolean} True if the node represents a function-like binding.
+ */
+export const isFunctionLikeBinding = node =>
+  isFunctionDeclaration( node ) ||
+  (
+    isVariableDeclarator( node ) &&
+    ( isFunctionExpression( node.init ) || isArrowFunctionExpression( node.init ) )
+  );

package/src/worker/webpack_loaders/workflow_rewriter/collect_target_imports.spec.js

@@ -12,24 +12,19 @@ function makeAst( source, filename ) {
 describe( 'collect_target_imports', () => {
   it( 'collects ESM imports for steps and workflows and flags changes', () => {
     const dir = mkdtempSync( join( tmpdir(), 'collect-esm-' ) );
-    writeFileSync( join( dir, 'steps.js' ), [
-      'export const StepA = step({ name: "step.a" })',
-      'export const StepB = step({ name: "step.b" })'
-    ].join( '\n' ) );
-    writeFileSync( join( dir, 'evaluators.js' ), [
-      'export const EvalA = evaluator({ name: "eval.a" })'
-    ].join( '\n' ) );
-    writeFileSync( join( dir, 'workflow.js' ), [
-      'export const FlowA = workflow({ name: "flow.a" })',
-      'export default workflow({ name: "flow.def" })'
-    ].join( '\n' ) );
+    writeFileSync( join( dir, 'steps.js' ), `
+export const StepA = step({ name: 'step.a' });
+export const StepB = step({ name: 'step.b' });` );
+    writeFileSync( join( dir, 'evaluators.js' ), 'export const EvalA = evaluator({ name: \'eval.a\' });' );
+    writeFileSync( join( dir, 'workflow.js' ), `
+export const FlowA = workflow({ name: 'flow.a' });
+export default workflow({ name: 'flow.def' });` );
 
-    const source = [
-      'import { StepA } from "./steps.js";',
-      'import { EvalA } from "./evaluators.js";',
-      'import WF, { FlowA } from "./workflow.js";',
-      'const x = 1;'
-    ].join( '\n' );
+    const source = `
+import { StepA } from './steps.js';
+import { EvalA } from './evaluators.js';
+import WF, { FlowA } from './workflow.js';
+const x = 1;`;
 
     const ast = makeAst( source, join( dir, 'file.js' ) );
     const { stepImports, evaluatorImports, flowImports } = collectTargetImports(
@@ -52,16 +47,15 @@ describe( 'collect_target_imports', () => {
 
   it( 'collects CJS requires and removes declarators (steps + default workflow)', () => {
     const dir = mkdtempSync( join( tmpdir(), 'collect-cjs-' ) );
-    writeFileSync( join( dir, 'steps.js' ), 'export const StepB = step({ name: "step.b" })\n' );
-    writeFileSync( join( dir, 'evaluators.js' ), 'export const EvalB = evaluator({ name: "eval.b" })\n' );
-    writeFileSync( join( dir, 'workflow.js' ), 'export default workflow({ name: "flow.c" })\n' );
+    writeFileSync( join( dir, 'steps.js' ), 'export const StepB = step({ name: \'step.b\' })' );
+    writeFileSync( join( dir, 'evaluators.js' ), 'export const EvalB = evaluator({ name: \'eval.b\' })' );
+    writeFileSync( join( dir, 'workflow.js' ), 'export default workflow({ name: \'flow.c\' })' );
 
-    const source = [
-      'const { StepB } = require("./steps.js");',
-      'const { EvalB } = require("./evaluators.js");',
-      'const WF = require("./workflow.js");',
-      'const obj = {};'
-    ].join( '\n' );
+    const source = `
+const { StepB } = require( './steps.js' );
+const { EvalB } = require( './evaluators.js' );
+const WF = require( './workflow.js' );
+const obj = {};`;
 
     const ast = makeAst( source, join( dir, 'file.js' ) );
     const { stepImports, evaluatorImports, flowImports } = collectTargetImports(

package/src/worker/webpack_loaders/workflow_rewriter/index.spec.js

@@ -19,25 +19,22 @@ function runLoader( source, resourcePath ) {
 describe( 'workflows_rewriter Webpack loader spec', () => {
   it( 'rewrites ESM imports and converts fn arrow to function', async () => {
     const dir = mkdtempSync( join( tmpdir(), 'ast-loader-esm-' ) );
-    writeFileSync( join( dir, 'steps.js' ), 'export const StepA = step({ name: \'step.a\' })\n' );
-    writeFileSync( join( dir, 'workflow.js' ), [
-      'export const FlowA = workflow({ name: \'flow.a\' })',
-      'export default workflow({ name: \'flow.def\' })'
-    ].join( '\n' ) );
-
-    const source = [
-      'import { StepA } from \'./steps.js\';',
-      'import FlowDef, { FlowA } from \'./workflow.js\';',
-      '',
-      'const obj = {',
-      '  fn: async (x) => {',
-      '    StepA(1);',
-      '    FlowA(2);',
-      '    FlowDef(3);',
-      '  }',
-      '}',
-      ''
-    ].join( '\n' );
+    writeFileSync( join( dir, 'steps.js' ), 'export const StepA = step({ name: \'step.a\' });' );
+    writeFileSync( join( dir, 'workflow.js' ), `
+export const FlowA = workflow({ name: 'flow.a' });
+export default workflow({ name: 'flow.def' });` );
+
+    const source = `
+import { StepA } from './steps.js';
+import FlowDef, { FlowA } from './workflow.js';
+
+const obj = {
+  fn: async (x) => {
+    StepA(1);
+    FlowA(2);
+    FlowDef(3);
+  }
+}`;
 
     const { code } = await runLoader( source, join( dir, 'file.js' ) );
 
@@ -53,18 +50,16 @@ describe( 'workflows_rewriter Webpack loader spec', () => {
 
   it( 'rewrites ESM shared_steps imports to invokeSharedStep', async () => {
     const dir = mkdtempSync( join( tmpdir(), 'ast-loader-esm-shared-' ) );
-    writeFileSync( join( dir, 'shared_steps.js' ), 'export const SharedA = step({ name: \'shared.a\' })\n' );
-
-    const source = [
-      'import { SharedA } from \'./shared_steps.js\';',
-      '',
-      'const obj = {',
-      '  fn: async (x) => {',
-      '    SharedA(1);',
-      '  }',
-      '}',
-      ''
-    ].join( '\n' );
+    writeFileSync( join( dir, 'shared_steps.js' ), 'export const SharedA = step({ name: \'shared.a\' });' );
+
+    const source = `
+import { SharedA } from './shared_steps.js';
+
+const obj = {
+  fn: async (x) => {
+    SharedA(1);
+  }
+}`;
 
     const { code } = await runLoader( source, join( dir, 'file.js' ) );
 
@@ -77,18 +72,16 @@ describe( 'workflows_rewriter Webpack loader spec', () => {
 
   it( 'rewrites CJS shared_steps requires to invokeSharedStep', async () => {
     const dir = mkdtempSync( join( tmpdir(), 'ast-loader-cjs-shared-' ) );
-    writeFileSync( join( dir, 'shared_steps.js' ), 'export const SharedB = step({ name: \'shared.b\' })\n' );
-
-    const source = [
-      'const { SharedB } = require(\'./shared_steps.js\');',
-      '',
-      'const obj = {',
-      '  fn: async (y) => {',
-      '    SharedB();',
-      '  }',
-      '}',
-      ''
-    ].join( '\n' );
+    writeFileSync( join( dir, 'shared_steps.js' ), 'export const SharedB = step({ name: \'shared.b\' });' );
+
+    const source = `
+const { SharedB } = require( './shared_steps.js' );
+
+const obj = {
+  fn: async (y) => {
+    SharedB();
+  }
+}`;
 
     const { code } = await runLoader( source, join( dir, 'file.js' ) );
 
@@ -101,21 +94,19 @@ describe( 'workflows_rewriter Webpack loader spec', () => {
 
   it( 'rewrites CJS requires and converts fn arrow to function', async () => {
     const dir = mkdtempSync( join( tmpdir(), 'ast-loader-cjs-' ) );
-    writeFileSync( join( dir, 'steps.js' ), 'export const StepB = step({ name: \'step.b\' })\n' );
-    writeFileSync( join( dir, 'workflow.js' ), 'export default workflow({ name: \'flow.c\' })\n' );
-
-    const source = [
-      'const { StepB } = require(\'./steps.js\');',
-      'const FlowDefault = require(\'./workflow.js\');',
-      '',
-      'const obj = {',
-      '  fn: async (y) => {',
-      '    StepB();',
-      '    FlowDefault();',
-      '  }',
-      '}',
-      ''
-    ].join( '\n' );
+    writeFileSync( join( dir, 'steps.js' ), 'export const StepB = step({ name: \'step.b\' });' );
+    writeFileSync( join( dir, 'workflow.js' ), 'export default workflow({ name: \'flow.c\' });' );
+
+    const source = `
+const { StepB } = require( './steps.js' );
+const FlowDefault = require( './workflow.js' );
+
+const obj = {
+  fn: async (y) => {
+    StepB();
+    FlowDefault();
+  }
+}`;
 
     const { code } = await runLoader( source, join( dir, 'file.js' ) );
 
@@ -130,22 +121,19 @@ describe( 'workflows_rewriter Webpack loader spec', () => {
 
   it( 'resolves top-level const name variables', async () => {
     const dir = mkdtempSync( join( tmpdir(), 'ast-loader-const-' ) );
-    writeFileSync( join( dir, 'steps.js' ), [
-      'const NAME = \'step.const\'',
-      'export const StepC = step({ name: NAME })'
-    ].join( '\n' ) );
-    writeFileSync( join( dir, 'workflow.js' ), [
-      'const WF = \'wf.const\'',
-      'export const FlowC = workflow({ name: WF })',
-      'const D = \'wf.def\'',
-      'export default workflow({ name: D })'
-    ].join( '\n' ) );
-
-    const source = [
-      'import { StepC } from \'./steps.js\';',
-      'import FlowDef, { FlowC } from \'./workflow.js\';',
-      'const obj = { fn: async () => { StepC(); FlowC(); FlowDef(); } }'
-    ].join( '\n' );
+    writeFileSync( join( dir, 'steps.js' ), `
+const NAME = 'step.const';
+export const StepC = step({ name: NAME });` );
+    writeFileSync( join( dir, 'workflow.js' ), `
+const WF = 'wf.const';
+export const FlowC = workflow({ name: WF });
+const D = 'wf.def';
+export default workflow({ name: D });` );
+
+    const source = `
+import { StepC } from './steps.js';
+import FlowDef, { FlowC } from './workflow.js';
+const obj = { fn: async () => { StepC(); FlowC(); FlowDef(); } }`;
 
     const { code } = await runLoader( source, join( dir, 'file.js' ) );
     expect( code ).toMatch( /this\.invokeStep\('step\.const'\)/ );
@@ -156,20 +144,17 @@ describe( 'workflows_rewriter Webpack loader spec', () => {
 
   it( 'throws on non-static name', async () => {
     const dir = mkdtempSync( join( tmpdir(), 'ast-loader-error-' ) );
-    writeFileSync( join( dir, 'steps.js' ), [
-      'function n() { return \'x\' }',
-      'export const StepX = step({ name: n() })'
-    ].join( '\n' ) );
-    writeFileSync( join( dir, 'workflow.js' ), [
-      'const base = \'a\'',
-      'export default workflow({ name: `${base}-b` })'
-    ].join( '\n' ) );
-
-    const source = [
-      'import { StepX } from \'./steps.js\';',
-      'import WF from \'./workflow.js\';',
-      'const obj = { fn: async () => { StepX(); WF(); } }'
-    ].join( '\n' );
+    writeFileSync( join( dir, 'steps.js' ), `
+function n() { return 'x'; }
+export const StepX = step({ name: n() });` );
+    writeFileSync( join( dir, 'workflow.js' ), `
+const base = 'a';
+export default workflow({ name: \`\${base}-b\` });` );
+
+    const source = `
+import { StepX } from './steps.js';
+import WF from './workflow.js';
+const obj = { fn: async () => { StepX(); WF(); } }`;
 
     await expect( runLoader( source, join( dir, 'file.js' ) ) ).rejects.toThrow( /Invalid (step|default workflow) name/ );
     rmSync( dir, { recursive: true, force: true } );

package/src/worker/webpack_loaders/workflow_rewriter/rewrite_fn_bodies.js

@@ -1,10 +1,147 @@
 import traverseModule from '@babel/traverse';
-import { isArrowFunctionExpression, isIdentifier, isFunctionExpression } from '@babel/types';
-import { toFunctionExpression, createThisMethodCall } from '../tools.js';
+import { isArrowFunctionExpression, isIdentifier } from '@babel/types';
+import {
+  toFunctionExpression,
+  createThisMethodCall,
+  isFunction,
+  bindThisAtCallSite,
+  isFunctionLikeBinding
+} from '../tools.js';
 
 // Handle CJS/ESM interop for Babel packages when executed as a webpack loader
 const traverse = traverseModule.default ?? traverseModule;
 
+/**
+ * Check whether a CallExpression callee is a simple Identifier.
+ * Only direct identifier calls are rewritten; member/dynamic calls are skipped.
+ *
+ * We only support rewriting `Foo()` calls that refer to imported steps/flows/evaluators
+ * or local call-chain functions. Calls like `obj.Foo()` or `(getFn())()` are out of scope.
+ *
+ * Examples:
+ * - Supported: `Foo()`
+ * - Skipped:   `obj.Foo()`, `(getFn())()`
+ *
+ * @param {import('@babel/traverse').NodePath} cPath - Path to a CallExpression node.
+ * @returns {boolean} True when callee is an Identifier.
+ */
+const isIdentifierCallee = cPath => isIdentifier( cPath.node.callee );
+
+/**
+ * Convert an ArrowFunctionExpression at the given path into a FunctionExpression
+ * to ensure dynamic `this` semantics inside the function body.
+ *
+ * Workflow code relies on `this` to invoke steps/flows (e.g., `this.invokeStep(...)`).
+ * Arrow functions capture `this` lexically, which would break that contract.
+ *
+ * If the node is an arrow, it is replaced by an equivalent FunctionExpression and
+ * the `state.rewrote` flag is set. If not an arrow, this is a no-op.
+ *
+ * @param {import('@babel/traverse').NodePath} nodePath - Path to a function node.
+ * @param {{ rewrote: boolean }} state - Mutation target to indicate a rewrite occurred.
+ * @returns {void}
+ */
+const normalizeArrowToFunctionPath = ( nodePath, state ) => {
+  if ( isArrowFunctionExpression( nodePath.node ) ) {
+    nodePath.replaceWith( toFunctionExpression( nodePath.node ) );
+    state.rewrote = true;
+  }
+};
+/**
+ * Rewrite calls inside a function body and collect call-chain functions discovered within.
+ * - Imported calls (steps/shared/evaluators/flows) are rewritten to `this.invokeX` or `this.startWorkflow`.
+ * - Local call-chain function calls are rewritten to `fn.call(this, ...)` to bind `this` correctly.
+ * - Returns a map of call-chain function name -> binding path for further recursive processing.
+ *
+ * @param {import('@babel/traverse').NodePath} bodyPath - Path to a function's body node.
+ * @param {Array<{ list: Array<any>, method: string, key: string }>} descriptors - Import rewrite descriptors.
+ * @param {{ rewrote: boolean }} state - Mutable state used to flag that edits were performed.
+ * @returns {Map<string, import('@babel/traverse').NodePath>} Discovered call-chain function bindings.
+ */
+const rewriteCallsInBody = ( bodyPath, descriptors, state ) => {
+  const callChainFunctions = new Map();
+  bodyPath.traverse( {
+    CallExpression: cPath => {
+      if ( !isIdentifierCallee( cPath ) ) {
+        return; // Only identifier callees are supported (skip member/dynamic)
+      }
+      const callee = cPath.node.callee;
+
+      // Rewrite imported calls (steps/shared/evaluators/flows)
+      for ( const { list, method, key } of descriptors ) {
+        const found = list.find( x => x.localName === callee.name );
+        if ( found ) {
+          const args = cPath.node.arguments;
+          cPath.replaceWith( createThisMethodCall( method, found[key], args ) );
+          state.rewrote = true;
+          return;
+        }
+      }
+
+      // Rewrite local call-chain function calls and track for recursive processing
+      const binding = cPath.scope.getBinding( callee.name );
+      if ( !binding ) {
+        return;
+      }
+      if ( !isFunctionLikeBinding( binding.path.node ) ) {
+        return; // Not a function-like binding
+      }
+
+      // Queue call-chain function for recursive processing
+      if ( !callChainFunctions.has( callee.name ) ) {
+        callChainFunctions.set( callee.name, binding.path );
+      }
+
+      // Bind `this` at callsite: fn(...) -> fn.call(this, ...)
+      cPath.replaceWith( bindThisAtCallSite( callee.name, cPath.node.arguments ) );
+      state.rewrote = true;
+    }
+  } );
+  return callChainFunctions;
+};
+
+/**
+ * Recursively process a call-chain function:
+ * - Ensures the function is a FunctionExpression (converts arrow when needed).
+ * - Rewrites calls inside the function using `rewriteCallsInBody`.
+ * - Follows nested call-chain functions depth-first while avoiding cycles via `processedFns`.
+ *
+ * @param {object} params - Params for processing a call-chain function.
+ * @param {string} params.name - Local identifier name in the current scope.
+ * @param {import('@babel/traverse').NodePath} params.bindingPath - Binding path of the function declaration.
+ * @param {{ rewrote: boolean }} params.state - Mutable state used to flag that edits were performed.
+ * @param {Array<{ list: Array<any>, method: string, key: string }>} params.descriptors - Import rewrite descriptors.
+ * @param {Set<string>} [params.processedFns] - Already processed names to avoid cycles.
+ */
+const processFunction = ( { name, bindingPath, state, descriptors, processedFns = new Set() } ) => {
+  // Avoid infinite loops for recursive/repeated references
+  if ( processedFns.has( name ) || bindingPath.removed ) {
+    return;
+  }
+  processedFns.add( name );
+
+  if ( bindingPath.isVariableDeclarator() ) {
+    // Case 1: const foo = <function or arrow>
+    const initPath = bindingPath.get( 'init' );
+    // Arrow functions capture `this` lexically; normalize for dynamic `this`
+    normalizeArrowToFunctionPath( initPath, state );
+    // Rewrite calls in body; collect nested call-chain functions from this scope
+    const callChainFunctions = rewriteCallsInBody( initPath.get( 'body' ), descriptors, state );
+    // DFS: process nested call-chain functions (processedFns prevents cycles)
+    callChainFunctions.forEach( ( childBindingPath, childName ) => {
+      processFunction( { name: childName, bindingPath: childBindingPath, state, descriptors, processedFns } );
+    } );
+  } else if ( bindingPath.isFunctionDeclaration() ) {
+    // Case 2: function foo(...) { ... }
+    // Function declarations already have dynamic `this`; no normalization needed
+    const callChainFunctions = rewriteCallsInBody( bindingPath.get( 'body' ), descriptors, state );
+    // Continue DFS into any functions called from this declaration
+    callChainFunctions.forEach( ( childBindingPath, childName ) => {
+      processFunction( { name: childName, bindingPath: childBindingPath, state, descriptors, processedFns } );
+    } );
+  }
+};
+
 /**
  * Rewrite calls to imported steps/workflows within `fn` object properties.
  * Converts arrow fns to functions and replaces `StepX(...)` with
@@ -21,49 +158,36 @@ const traverse = traverseModule.default ?? traverseModule;
  */
 export default function rewriteFnBodies( { ast, stepImports, sharedStepImports = [], evaluatorImports, flowImports } ) {
   const state = { rewrote: false };
+  // Build rewrite descriptors once per traversal
+  const descriptors = [
+    { list: stepImports, method: 'invokeStep', key: 'stepName' },
+    { list: sharedStepImports, method: 'invokeSharedStep', key: 'stepName' },
+    { list: evaluatorImports, method: 'invokeEvaluator', key: 'evaluatorName' },
+    { list: flowImports, method: 'startWorkflow', key: 'workflowName' }
+  ];
   traverse( ast, {
     ObjectProperty: path => {
-      // If it is not fn property: skip
+      // Only transform object properties named 'fn'
       if ( !isIdentifier( path.node.key, { name: 'fn' } ) ) {
         return;
       }
 
       const val = path.node.value;
 
-      // if it is not function, return
-      if ( !isFunctionExpression( val ) && !isArrowFunctionExpression( val ) ) {
+      // Only functions (including arrows) are eligible
+      if ( !isFunction( val ) && !isArrowFunctionExpression( val ) ) {
         return;
       }
 
-      // replace arrow fn in favor of a function
-      if ( isArrowFunctionExpression( val ) ) {
-        const func = toFunctionExpression( val );
-        path.get( 'value' ).replaceWith( func );
-        state.rewrote = true;
-      }
+      // Normalize arrow to function for correct dynamic `this`
+      normalizeArrowToFunctionPath( path.get( 'value' ), state );
 
-      path.get( 'value.body' ).traverse( {
-        CallExpression: cPath => {
-          const callee = cPath.node.callee;
-          if ( !isIdentifier( callee ) ) {
-            return;
-          } // Skip: complex callee not supported
-          const descriptors = [
-            { list: stepImports, method: 'invokeStep', key: 'stepName' },
-            { list: sharedStepImports, method: 'invokeSharedStep', key: 'stepName' },
-            { list: evaluatorImports, method: 'invokeEvaluator', key: 'evaluatorName' },
-            { list: flowImports, method: 'startWorkflow', key: 'workflowName' }
-          ];
-          for ( const { list, method, key } of descriptors ) {
-            const found = list.find( x => x.localName === callee.name );
-            if ( found ) {
-              const args = cPath.node.arguments;
-              cPath.replaceWith( createThisMethodCall( method, found[key], args ) );
-              state.rewrote = true;
-              return;
-            }
-          }
-        }
+      // Rewrite the main workflow fn body and collect call-chain functions discovered within it
+      const callChainFunctions = rewriteCallsInBody( path.get( 'value.body' ), descriptors, state );
+
+      // Recursively rewrite call-chain functions and any functions they call
+      callChainFunctions.forEach( ( bindingPath, name ) => {
+        processFunction( { name, bindingPath, state, descriptors } );
       } );
     }
   } );