@output.ai/core 0.1.8 → 0.1.9

package/README.md

@@ -1,246 +1,98 @@
-# Core
+# @output.ai/core
 
-Provides tools to develop and run a workflow, which is a well defined logical unit of work.
+Workflow orchestration and worker runtime for building durable LLM applications with Temporal.
 
-## Structure
+[![npm version](https://img.shields.io/npm/v/@output.ai/core)](https://www.npmjs.com/package/@output.ai/core)
+[![Documentation](https://img.shields.io/badge/docs-docs.output.ai-blue)](https://docs.output.ai/packages/core)
 
-Workflows are defined using core functions ("workflow", "step", "evaluator"), these are defined in separate files, and must be placed within the same folder:
-
-```
-└ workflows
-  └ example
-    ├ workflow.ts|js <- workflow entry point
-    ├ steps.ts|js <- file containing steps used by the workflow
-    ├ evaluators.ts|js <- file containing evaluating functions
-    └ prompt.prompt <- a prompt file
-  └ other-example
+## Installation
 
+```bash
+npm install @output.ai/core
 ```
 
-Workflows are the orchestrator and steps are executors. So the workflow only call the steps and the steps call the IO operations, like APIs, DBs, LLMs, etc. Evaluators are just another different flavor for steps, they work the same, but must return an `EvaluationResult` object.
-
-## Components
+## Quick Start
 
-### Workflow
+```typescript
+// workflow.ts
+import { workflow, z } from '@output.ai/core';
+import { processData } from './steps.js';
 
-The main code, must contain only deterministic orchestration code.
-
-File: `workflow.js`
-
-Example:
-```js
-import { workflow, z } from '@output.ai/workflow';
-import { guessByName } from './steps.js';
-
-export default workflow( {
-  name: 'guessMyProfession',
-  description: 'Guess a person profession by its name',
-  inputSchema: z.object( {
-    name: z.string()
-  } ),
-  outputSchema: z.object( {
-    profession: z.string()
-  } ),
-  fn: async input => {
-    const profession = await guessByName( input.name );
-    return { profession };
+export default workflow({
+  name: 'myWorkflow',
+  inputSchema: z.object({ text: z.string() }),
+  outputSchema: z.object({ result: z.string() }),
+  fn: async (input) => {
+    const result = await processData(input.text);
+    return { result };
   }
-})
+});
 ```
 
-Workflows can only import the following files
-
-#### Components
-- `evaluators.js`
-- `shared_steps.js`
-- `steps.js`
-- `workflow.js`
+```typescript
+// steps.ts
+import { step, z } from '@output.ai/core';
 
-#### Core library
-- `@output.ai.core`
-
-#### Whitelisted files
-- `types.js`
-- `consts.js`
-- `constants.js`
-- `vars.js`
-- `variables.js`
-- `utils.js`
-- `tools.js`
-- `functions.js`
-- `shared.js`
-
-### Step
-
-Re-usable units of work that can contain IO, used by the workflow.
-
-File: `steps.js`
-
-Example:
-```js
-import { api } from './api.js'
-
-export const guessByName = step( {
-  name: 'guessByName',
+export const processData = step({
+  name: 'processData',
   inputSchema: z.string(),
   outputSchema: z.string(),
-  fn: async name => {
-    const res = await api.consumer( name );
-    return res.body;
-  }
-} )
-```
-
-### Shared Steps
-
-By default, steps are exclusive to the workflow, so it is not passible to use these steps from elsewhere. In order to have shared steps and make them accessible in different workflows, create a shared steps file. This file can be relatively imported anywhere.
-
-File: `shared_steps.js`
-
-Example:
-```js
-export const mySharedStep = step( {
-  name: 'mySharedStep',
-  ...
-} )
-```
-
-And the usage is the same as any step:
-`workflow.js`
-```js
-import { mySharedStep } from '../../tools/shared_steps.js'
-```
-
-### Evaluators
-
-Steps that analyze LLM response, or take other measurements are contained in evaluators.
-
-File: `evaluators.js`
-
-Example:
-```js
-import { evaluator, EvaluationStringResult } from './api.js'
-
-export const judgeResult = evaluator( {
-  name: 'judgeResult',
-  inputSchema: z.string(),
-  fn: async name => {
-    ...
-    return new EvaluationStringResult({
-      value: 'good',
-      confidence: .95
-    });
+  fn: async (text) => {
+    return text.toUpperCase();
   }
-} )
+});
 ```
 
-Its usage is the same as steps:
-`workflow.js`
-```js
-import { workflow, z } from '@output.ai/workflow';
-import { judgeResult } from './evaluators.js';
-
-export default workflow( {
-  name: 'guessMyProfession',
-  inputSchema: z.object( {
-    name: z.string()
-  } ),
-  outputSchema: z.object( {
-    result: z.string()
-  } ),
-  fn: async input => {
-    const judgment = await judgeResult( input.name );
-    return { result: judgement.value };
-  }
-})
-```
-
-## Webhooks
-
-Workflows can call webhooks that will stop their execution until an answer is given back.
-
-```js
-import { workflow, createWebhook } from '@output.ai/workflow';
-import { guessByName } from './steps.js';
-
-export default workflow( {
-  ...
-  fn: async input => {
-    ...
-
-    const result = await createWebhook( {
-      url: 'http://xxx.xxx/feedback',
-      payload: {
-        progressSoFar: 'plenty'
-      }
-    } );
-
-  }
-})
-```
-
-The url of the example will receive the payload, plus the workflowId:
-
-```js
-{
-  workflowId: '', // alphanumerical id of the workflow execution,
-  payload: { }, // the payload sent using tools.webhook()
-}
-```
-
-To resume the workflow, a POST has to be made with a response payload and the workflowId.
-
-- Production: `https://output-api-production.onrender.com/workflow/feedback`
-- Local: `http://localhost:3001/workflow/feedback`
-
-Example:
-
-```bash
-POST http://locahost:3001/workflow/feedback
-  {
-    workflowId,
-    payload: {}
-  }
-```
-
-## Options
-
-All core interface functions: workflow, step, evaluator have similar signature, with the following options:
-- name: The function name, used to call it internally and identify it in the trace files, must be a code friendly string;
-- description: Human description of the workflow/step, used for the catalog;
-- inputSchema: a zod object indicating the type of the argument received by the `fn` function. It is validated. Omit if it doesn't have input arguments;
-- outputSchema: a zod object indicating the type of that the `fn` function returns. It is validated. Omit if it is void. Evaluators do not have this option, since they must always return an EvaluationResult object;
-- fn: The actual implementation of the workflow/step, including all its logic.
-- options: Advanced options that will overwrite Temporal's ActivityOptions when calling activities.
-
-  If used on `workflow()` it will apply for all activities. If used on `step()` or `evaluator()` it will apply only to that underlying activity. If changed in both places, the end value will be a merge between the initial values, workflow values and the step values.
-
-  Order of precedence
-  `step options > workflow options > default options`
-
-## Developing
-
-To develop workflows you need the code, which will be called the worker, the API and the engine (Temporal).
-
-After having the API and the engine running, to start the worker just run:
-
-```js
-`npm run outputai`
+## Key Exports
+
+| Export | Description |
+|--------|-------------|
+| `workflow` | Define orchestration logic that coordinates steps |
+| `step` | Define reusable units of work that handle I/O |
+| `evaluator` | Define steps that return evaluation results |
+| `createWebhook` | Pause workflow execution until external input |
+| `z` | Zod schema library for input/output validation |
+
+## File Structure
+
+Each workflow lives in its own directory:
+
+```text
+src/workflows/
+└── my-workflow/
+    ├── workflow.ts           # Workflow definition
+    ├── steps.ts              # Step implementations
+    ├── evaluators.ts         # Evaluators (optional)
+    ├── prompts/              # LLM prompt templates
+    │   └── prompt@v1.prompt
+    └── scenarios/            # Test scenarios
+        └── test_input.json
 ```
 
-## Env variables
-
-Necessary env variables to run the worker locally:
-
-- `TEMPORAL_ADDRESS`: The temporal backend address, prefer the remote;
-- `TEMPORAL_NAMESPACE`: The name of the namespace, if using remote, use: "output-production.i0jzq";
-- `TEMPORAL_API_KEY`: The API key to access remote temporal. If using local temporal, leave it blank;
-- `CATALOG_ID`: The name of the local catalog, always set this. Use your email;
-- `API_AUTH_KEY`: The API key to access the Framework API. Local can be blank, remote use the proper API Key;
-- `TRACE_LOCAL_ON`: A "stringbool" value indicating if traces should be saved locally, needs REDIS_URL;
-- `TRACE_REMOTE_ON`: A "stringbool" value indicating if traces should be saved remotely, needs REDIS_URL and AWS_* secrets;
-- `REDIS_URL`: The redis address to connect. Only necessary when any type of trace is enabled;
-- `TRACE_REMOTE_S3_BUCKET`: The AWS S3 bucket to send the traces. Only necessary when remote trace is enabled;
-- `AWS_REGION`: AWS region to connect to send the traces, must match the bucket region. Only necessary when remote trace is enabled;
-- `AWS_ACCESS_KEY_ID`: AWS key id. Only necessary when remote trace is enabled;
-- `AWS_SECRET_ACCESS_KEY`: AWS secrete. Only necessary when remote trace is enabled;
+## Environment Variables
+
+The worker reads these environment variables:
+
+| Variable | Description |
+|----------|-------------|
+| `TEMPORAL_ADDRESS` | Temporal backend address |
+| `TEMPORAL_NAMESPACE` | Temporal namespace name |
+| `TEMPORAL_API_KEY` | API key for remote Temporal (leave blank for local) |
+| `CATALOG_ID` | **Required.** Name of the local catalog (use your email) |
+| `API_AUTH_KEY` | API key for Framework API (blank for local, required for remote) |
+| `TRACE_LOCAL_ON` | Enable local trace saving (requires `REDIS_URL`) |
+| `TRACE_REMOTE_ON` | Enable remote trace saving (requires `REDIS_URL` and AWS secrets) |
+| `REDIS_URL` | Redis address (required when tracing is enabled) |
+| `TRACE_REMOTE_S3_BUCKET` | AWS S3 bucket for traces (required for remote tracing) |
+| `AWS_REGION` | AWS region matching the S3 bucket (required for remote tracing) |
+| `AWS_ACCESS_KEY_ID` | AWS key ID (required for remote tracing) |
+| `AWS_SECRET_ACCESS_KEY` | AWS secret key (required for remote tracing) |
+
+## Documentation
+
+For comprehensive documentation, visit:
+
+- [Package Reference](https://docs.output.ai/packages/core)
+- [Workflows Guide](https://docs.output.ai/core/workflows)
+- [Steps Guide](https://docs.output.ai/core/steps)
+- [Getting Started](https://docs.output.ai/quickstart)

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@output.ai/core",
-  "version": "0.1.8",
+  "version": "0.1.9",
   "description": "The core module of the output framework",
   "type": "module",
   "exports": {
@@ -40,7 +40,10 @@
     "stacktrace-parser": "0.1.11",
     "zod": "4.1.12"
   },
-  "license": "UNLICENSED",
+  "license": "Apache-2.0",
+  "publishConfig": {
+    "access": "public"
+  },
   "imports": {
     "#consts": "./src/consts.js",
     "#errors": "./src/errors.js",

package/src/tracing/trace_engine.js

@@ -9,11 +9,13 @@ const traceBus = new EventEmitter();
 const processors = [
   {
     isOn: isStringboolTrue( process.env.TRACE_LOCAL_ON ),
+    name: 'LOCAL',
     init: localProcessor.init,
     exec: localProcessor.exec
   },
   {
     isOn: isStringboolTrue( process.env.TRACE_REMOTE_ON ),
+    name: 'REMOTE',
     init: s3Processor.init,
     exec: s3Processor.exec
   }
@@ -25,7 +27,13 @@ const processors = [
 export const init = async () => {
   for ( const p of processors.filter( p => p.isOn ) ) {
     await p.init();
-    traceBus.addListener( 'entry', p.exec );
+    traceBus.addListener( 'entry', async ( ...args ) => {
+      try {
+        await p.exec( ...args );
+      } catch ( error ) {
+        console.error( `[Tracing] "${p.name}" processor execution error.`, error );
+      }
+    } );
   }
 };
 

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 } );
       } );
     }
   } );

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

@@ -1,17 +1,19 @@
 import { describe, it, expect } from 'vitest';
+import generatorModule from '@babel/generator';
 import { parse } from '../tools.js';
 import rewriteFnBodies from './rewrite_fn_bodies.js';
 
+const generate = generatorModule.default ?? generatorModule;
+
 describe( 'rewrite_fn_bodies', () => {
   it( 'converts arrow to function and rewrites step/workflow calls', () => {
-    const src = [
-      'const obj = {',
-      '  fn: async (x) => {',
-      '    StepA(1);',
-      '    FlowB(2);',
-      '  }',
-      '}'
-    ].join( '\n' );
+    const src = `
+const obj = {
+  fn: async x => {
+    StepA( 1 );
+    FlowB( 2 );
+  }
+}`;
     const ast = parse( src, 'file.js' );
     const stepImports = [ { localName: 'StepA', stepName: 'step.a' } ];
     const flowImports = [ { localName: 'FlowB', workflowName: 'flow.b' } ];
@@ -24,13 +26,12 @@ describe( 'rewrite_fn_bodies', () => {
   } );
 
   it( 'rewrites evaluator calls to this.invokeEvaluator', () => {
-    const src = [
-      'const obj = {',
-      '  fn: async (x) => {',
-      '    EvalA(3);',
-      '  }',
-      '}'
-    ].join( '\n' );
+    const src = `
+const obj = {
+  fn: async x => {
+    EvalA(3);
+  }
+};`;
     const ast = parse( src, 'file.js' );
     const evaluatorImports = [ { localName: 'EvalA', evaluatorName: 'eval.a' } ];
     const rewrote = rewriteFnBodies( { ast, stepImports: [], evaluatorImports, flowImports: [] } );
@@ -43,5 +44,80 @@ describe( 'rewrite_fn_bodies', () => {
     const rewrote = rewriteFnBodies( { ast, stepImports: [], evaluatorImports: [], flowImports: [] } );
     expect( rewrote ).toBe( false );
   } );
+
+  it( 'rewrites helper calls and helper bodies (steps and evaluators)', () => {
+    const src = `
+const foo = async () => {
+  StepA( 1 );
+};
+
+function bar( x ) {
+  EvalA( x );
+}
+
+const obj = {
+  fn: async x => {
+    foo();
+    bar( 2 );
+  }
+}`;
+
+    const ast = parse( src, 'file.js' );
+    const stepImports = [ { localName: 'StepA', stepName: 'step.a' } ];
+    const evaluatorImports = [ { localName: 'EvalA', evaluatorName: 'eval.a' } ];
+
+    const rewrote = rewriteFnBodies( { ast, stepImports, sharedStepImports: [], evaluatorImports, flowImports: [] } );
+    expect( rewrote ).toBe( true );
+
+    const { code } = generate( ast, { quotes: 'single' } );
+
+    // Helper calls in fn are rewritten to call(this, ...)
+    expect( code ).toMatch( /foo\.call\(this\)/ );
+    expect( code ).toMatch( /bar\.call\(this,\s*2\)/ );
+
+    // Inside helpers, calls are rewritten
+    expect( code ).toMatch( /this\.invokeStep\(([\"'])step\.a\1,\s*1\)/ );
+    expect( code ).toMatch( /this\.invokeEvaluator\(([\"'])eval\.a\1,\s*x\)/ );
+
+    // Arrow helper converted to function expression to allow dynamic this
+    expect( code ).toMatch( /const foo = async function/ );
+  } );
+
+  it( 'rewrites nested helper chains until the step invocation', () => {
+    const src = `
+const foo = () => {
+  bar();
+};
+
+function bar() {
+  baz( 42 );
+}
+
+const baz = n => {
+  StepA( n );
+};
+
+const obj = {
+  fn: async () => {
+    foo();
+  }
+}`;
+
+    const ast = parse( src, 'file.js' );
+    const stepImports = [ { localName: 'StepA', stepName: 'step.a' } ];
+    const rewrote = rewriteFnBodies( { ast, stepImports, evaluatorImports: [], flowImports: [] } );
+    expect( rewrote ).toBe( true );
+
+    const { code } = generate( ast, { quotes: 'single' } );
+    // Calls along the chain are bound with this
+    expect( code ).toMatch( /foo\.call\(this\)/ );
+    expect( code ).toMatch( /bar\.call\(this\)/ );
+    expect( code ).toMatch( /baz\.call\(this,\s*42\)/ );
+    // Deep step rewrite in the last helper
+    expect( code ).toMatch( /this\.invokeStep\(([\"'])step\.a\1,\s*n\)/ );
+    // Arrow helpers converted to functions
+    expect( code ).toMatch( /const foo = function/ );
+    expect( code ).toMatch( /const baz = function/ );
+  } );
 } );