Support projects with Node.js step dependencies in vitest plugin (#1524)

Co-authored-by: Peter Wielander <mittgfu@gmail.com>

Author: matchai

.changeset/fix-builder-ts-extensions.md

@@ -0,0 +1,5 @@
+---
+"@workflow/builders": patch
+---
+
+Fix dependency resolution for step imports with .ts, .mts, and .cts extensions

.changeset/fix-vitest-node-deps.md

@@ -0,0 +1,5 @@
+---
+"@workflow/vitest": patch
+---
+
+Fix step dependencies not being mockable when imported from TypeScript files

docs/content/docs/testing/index.mdx

@@ -85,7 +85,7 @@ Unit testing works well for individual steps. A simple workflow that only calls
 For workflows that rely on runtime features like [hooks](/docs/foundations/hooks), [webhooks](/docs/foundations/hooks#understanding-webhooks), [`sleep()`](/docs/api-reference/workflow/sleep), or error retries, you need to test against a real workflow setup. The `@workflow/vitest` plugin handles everything automatically — it compiles your workflow directives, builds the runtime bundles, and executes workflows entirely in-process. No server required.
 
 <Callout type="warn">
-Inside integration tests, which run the full workflow runtime, `vi.mock()` and related calls do not work — neither for your own modules nor for third-party npm packages. All step dependencies are inlined into the compiled bundle by esbuild, bypassing Vitest's module system entirely. To test steps with mocked dependencies, use [unit tests](#unit-testing-steps) instead. Consider dependency injection or environment variable-based conditional logic for controlling behavior in integration tests.
+`vi.mock()` and related calls do _not_ work inside workflow functions, only step functions. Your workflow functions cannot import third party code that needs to be mocked. Mocking works for npm packages imported in step functions. If something needs to be mocked, it likely belongs inside a step function either way.
 </Callout>
 
 ### Vitest Configuration

packages/builders/src/base-builder.ts

@@ -110,6 +110,7 @@ export abstract class BaseBuilder {
       '**/.output/**',
       '**/.vercel/**',
       '**/.workflow-data/**',
+      '**/.workflow-vitest/**',
       '**/.well-known/workflow/**',
       '**/.svelte-kit/**',
       '**/.turbo/**',

packages/builders/src/get-input-files.test.ts

@@ -101,6 +101,7 @@ describe('getInputFiles', () => {
     writeFile(srcDir, '.vercel/output/step.ts');
     writeFile(srcDir, '.svelte-kit/output/step.ts');
     writeFile(srcDir, '.workflow-data/state.ts');
+    writeFile(srcDir, '.workflow-vitest/workflows.mjs');
     writeFile(srcDir, '.well-known/workflow/route.ts');
     writeFile(srcDir, '.turbo/cache/build.ts');
     writeFile(srcDir, '.cache/babel/plugin.js');
@@ -131,6 +132,9 @@ describe('getInputFiles', () => {
     expect(files).not.toContain(
       normalize(join(srcDir, '.workflow-data/state.ts'))
     );
+    expect(files).not.toContain(
+      normalize(join(srcDir, '.workflow-vitest/workflows.mjs'))
+    );
     expect(files).not.toContain(
       normalize(join(srcDir, '.well-known/workflow/route.ts'))
     );

packages/builders/src/swc-esbuild-plugin.test.ts

@@ -27,6 +27,64 @@ function writeFile(path: string, contents = ''): void {
   writeFileSync(path, contents, 'utf-8');
 }
 
+describe('createSwcPlugin externalizeNonSteps', () => {
+  let testRoot: string;
+
+  beforeEach(() => {
+    testRoot = mkdtempSync(join(realTmpdir, 'workflow-swc-plugin-'));
+    applySwcTransformMock.mockReset();
+    applySwcTransformMock.mockImplementation(
+      async (_filename: string, source: string) => ({
+        code: source,
+        workflowManifest: {},
+      })
+    );
+  });
+
+  afterEach(() => {
+    rmSync(testRoot, { recursive: true, force: true });
+  });
+
+  it.each([
+    { inputExt: '.ts', outputExt: '.js' },
+    { inputExt: '.tsx', outputExt: '.js' },
+    { inputExt: '.mts', outputExt: '.mjs' },
+    { inputExt: '.cts', outputExt: '.cjs' },
+  ])('rewrites externalized $inputExt imports to $outputExt', async ({
+    inputExt,
+    outputExt,
+  }) => {
+    const outdir = join(testRoot, 'out');
+    const srcDir = join(testRoot, 'src');
+    const stepFile = join(srcDir, 'step.ts');
+
+    writeFile(join(srcDir, `dep${inputExt}`), 'export const dep = {};');
+    writeFile(stepFile, `import { dep } from './dep';\nconsole.log(dep);`);
+
+    const result = await esbuild.build({
+      entryPoints: [stepFile],
+      absWorkingDir: testRoot,
+      outdir,
+      bundle: true,
+      format: 'esm',
+      platform: 'node',
+      write: false,
+      plugins: [
+        createSwcPlugin({
+          mode: 'step',
+          entriesToBundle: [stepFile],
+          outdir,
+        }),
+      ],
+    });
+
+    expect(result.errors).toHaveLength(0);
+    const output = result.outputFiles[0].text;
+    expect(output).toContain(`/dep${outputExt}`);
+    expect(output).not.toContain(`/dep${inputExt}`);
+  });
+});
+
 describe('createSwcPlugin projectRoot', () => {
   let testRoot: string;
 

packages/builders/src/swc-esbuild-plugin.ts

@@ -124,15 +124,24 @@ export function createSwcPlugin(options: SwcPluginOptions): Plugin {
           const isFilePath =
             args.path.startsWith('.') || args.path.startsWith('/');
 
-          return {
-            external: true,
-            path: isFilePath
-              ? relative(options.outdir || process.cwd(), resolvedPath).replace(
-                  /\\/g,
-                  '/'
-                )
-              : args.path,
-          };
+          let externalPath: string;
+          if (isFilePath) {
+            externalPath = relative(
+              options.outdir || process.cwd(),
+              resolvedPath
+            ).replace(/\\/g, '/');
+
+            // Rewrite TypeScript extensions to their JS equivalents so the
+            // externalized import is loadable by Node's native ESM loader.
+            externalPath = externalPath
+              .replace(/\.tsx?$/, '.js')
+              .replace(/\.mts$/, '.mjs')
+              .replace(/\.cts$/, '.cjs');
+          } else {
+            externalPath = args.path;
+          }
+
+          return { external: true, path: externalPath };
         } catch (_) {}
         return null;
       });

packages/vitest/src/index.ts

@@ -86,7 +86,9 @@ function getOutDir(cwd: string): string {
 export function workflow(): Plugin[] {
   const dir = fileURLToPath(new URL('.', import.meta.url));
   return [
-    workflowTransformPlugin(),
+    workflowTransformPlugin({
+      exclude: [join(process.cwd(), '.workflow-vitest') + '/'],
+    }),
     {
       name: 'workflow:vitest',
       config() {
@@ -139,17 +141,28 @@ export async function setupWorkflowTests(
   const cwd = options?.cwd ?? process.cwd();
   const outDir = getOutDir(cwd);
 
-  const workflowsModule = await import(
-    /* @vite-ignore */ pathToFileURL(join(outDir, 'workflows.mjs')).href
-  );
-  const stepsModule = await import(
-    /* @vite-ignore */ pathToFileURL(join(outDir, 'steps.mjs')).href
-  );
-
-  const workflowHandler = workflowsModule.POST as (
-    req: Request
-  ) => Promise<Response>;
-  const stepHandler = stepsModule.POST as (req: Request) => Promise<Response>;
+  // Lazy-load bundles on first dispatch instead of eagerly at setup time.
+  // Eager native import() during setupFiles loads step dependencies into
+  // the module cache before vi.mock() can intercept them, breaking mocks
+  // in unit tests that never execute workflows.
+  function createLazyHandler(
+    bundlePath: string
+  ): (req: Request) => Promise<Response> {
+    let handler: ((req: Request) => Promise<Response>) | undefined;
+    let loading: Promise<(req: Request) => Promise<Response>> | undefined;
+
+    return async (req: Request) => {
+      if (!handler) {
+        // If the import rejects (e.g. missing bundle), the rejected promise is
+        // cached so all subsequent calls fail fast with the same error.
+        loading ??= import(
+          /* @vite-ignore */ pathToFileURL(bundlePath).href
+        ).then((mod) => mod.POST as (req: Request) => Promise<Response>);
+        handler = await loading;
+      }
+      return handler(req);
+    };
+  }
 
   // Each vitest worker uses a unique tag to isolate its test data.
   // All workers write to the shared .workflow-data directory so runs
@@ -163,8 +176,14 @@ export async function setupWorkflowTests(
   await world.start?.();
   await world.clear();
 
-  world.registerHandler('__wkf_workflow_', workflowHandler);
-  world.registerHandler('__wkf_step_', stepHandler);
+  world.registerHandler(
+    '__wkf_workflow_',
+    createLazyHandler(join(outDir, 'workflows.mjs'))
+  );
+  world.registerHandler(
+    '__wkf_step_',
+    createLazyHandler(join(outDir, 'steps.mjs'))
+  );
 
   setWorld(world);
 }

workbench/vitest/test/mock.test.ts

@@ -1,26 +1,48 @@
 /**
  * Test whether vi.mock() works for third-party npm packages
  * in the vitest integration test environment.
- *
- * This test verifies that mocking third-party packages (like `ms`)
- * does NOT work in integration tests because the step bundle is compiled
- * by esbuild which inlines all dependencies. vi.mock() only affects
- * the vitest module graph, not the pre-compiled step bundle.
  */
+
+import ms from 'ms';
 import { describe, expect, it, vi } from 'vitest';
 import { start } from 'workflow/api';
-import { durationWorkflow } from '../workflows/third-party.js';
+import {
+  durationWorkflow,
+  durationWorkflowInline,
+  durationWorkflowStepUtil,
+} from '../workflows/third-party.js';
 
 vi.mock('ms', () => ({
   default: () => 42,
 }));
 
 describe('third-party mocking', () => {
-  it('vi.mock does NOT intercept third-party imports in steps', async () => {
+  it('vi.mock intercepts external npm package used in step', async () => {
+    // Mock works outside the workflow bundle
+    expect(ms('1h')).toBe(42);
+
     const run = await start(durationWorkflow, ['1h']);
     const result = await run.returnValue;
-    // If the mock worked, result would be { ms: 42 }
-    // Since the step bundle inlines dependencies, the real ms() is used
-    expect(result).toEqual({ ms: 3_600_000 }); // real ms('1h') = 3600000
+
+    // Mock works inside the workflow bundle
+    expect(result).toEqual({ ms: 42 });
+  });
+
+  it.fails('vi.mock intercepts external npm package used in workflow', async () => {
+    expect(ms('1h')).toBe(42);
+
+    const run = await start(durationWorkflowInline, ['1h']);
+    const result = await run.returnValue;
+
+    // Mock doesn't yet work inside the workflow bundle
+    expect(result).toEqual({ ms: 42 });
+  });
+
+  it('vi.mock intercepts internal import used in step', async () => {
+    const run = await start(durationWorkflowStepUtil, ['1h']);
+    const result = await run.returnValue;
+
+    // Mock doesn't work for internalized local dependencies
+    expect(result).toEqual({ ms: 42 });
   });
 });

workbench/vitest/workflows/third-party.ts

@@ -3,14 +3,30 @@
  * Used to test whether vi.mock() works for third-party dependencies.
  */
 import ms from 'ms';
+import { formatDurationUtil } from './utils';
 
-async function formatDuration(duration: string) {
+export async function formatDurationStepUsingExternal(duration: string) {
+  'use step';
+  return formatDurationUtil(duration);
+}
+
+export async function formatDurationStep(duration: string) {
   'use step';
   return ms(duration);
 }
 
 export async function durationWorkflow(duration: string) {
   'use workflow';
-  const result = await formatDuration(duration);
+  const result = await formatDurationStep(duration);
   return { ms: result };
 }
+
+export async function durationWorkflowInline(duration: string) {
+  'use workflow';
+  return { ms: ms(duration) };
+}
+
+export async function durationWorkflowStepUtil(duration: string) {
+  'use workflow';
+  return { ms: await formatDurationStepUsingExternal(duration) };
+}