[core] Enforce single ALS context to protect against duplicate module and caching issues (#1591)

Author: VaguelySerious

.changeset/fix-step-context-singleton.md

@@ -0,0 +1,5 @@
+---
+"@workflow/core": patch
+---
+
+Fix step `contextStorage` global _potentially_ seeing dual-instance issues when bundlers create multiple copies of the module.

packages/core/e2e/e2e.test.ts

@@ -2156,4 +2156,20 @@ describe('e2e', () => {
       });
     }
   );
+
+  test(
+    'metadataFromHelperWorkflow - getWorkflowMetadata/getStepMetadata work from module-level helper (#1577)',
+    { timeout: 60_000 },
+    async () => {
+      const run = await start(await e2e('metadataFromHelperWorkflow'), [
+        'smoke-test',
+      ]);
+      const returnValue = await run.returnValue;
+
+      expect(returnValue.label).toBe('smoke-test');
+      expect(typeof returnValue.workflowRunId).toBe('string');
+      expect(typeof returnValue.stepId).toBe('string');
+      expect(returnValue.attempt).toBeGreaterThanOrEqual(1);
+    }
+  );
 });

packages/core/src/step/context-storage.test.ts

@@ -0,0 +1,40 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { describe, expect, it } from 'vitest';
+
+const CONTEXT_STORAGE_SYMBOL = Symbol.for('WORKFLOW_STEP_CONTEXT_STORAGE');
+
+describe('contextStorage singleton', () => {
+  it('returns the same AsyncLocalStorage instance across multiple imports', async () => {
+    // Import the module twice (simulating what bundlers might do)
+    const mod1 = await import('./context-storage.js');
+    const mod2 = await import('./context-storage.js');
+
+    expect(mod1.contextStorage).toBe(mod2.contextStorage);
+  });
+
+  it('shares the same instance stored on globalThis via Symbol.for()', async () => {
+    const { contextStorage } = await import('./context-storage.js');
+    const globalInstance = (globalThis as any)[CONTEXT_STORAGE_SYMBOL];
+
+    expect(globalInstance).toBe(contextStorage);
+    expect(globalInstance).toBeInstanceOf(AsyncLocalStorage);
+  });
+
+  it('preserves context from run() when getStore() is called from a separately-constructed reference', async () => {
+    // This simulates the dual-module-instance problem:
+    // step-handler sets context via one reference, user code reads via another.
+    // With the Symbol.for() singleton fix, both references point to the same instance.
+    const { contextStorage } = await import('./context-storage.js');
+    const globalInstance = (globalThis as any)[
+      CONTEXT_STORAGE_SYMBOL
+    ] as AsyncLocalStorage<{ value: string }>;
+
+    let storeFromGlobal: { value: string } | undefined;
+
+    globalInstance.run({ value: 'test-context' }, () => {
+      storeFromGlobal = contextStorage.getStore() as any;
+    });
+
+    expect(storeFromGlobal).toEqual({ value: 'test-context' });
+  });
+});

packages/core/src/step/context-storage.ts

@@ -3,10 +3,35 @@ import type { CryptoKey } from '../encryption.js';
 import type { WorkflowMetadata } from '../workflow/get-workflow-metadata.js';
 import type { StepMetadata } from './get-step-metadata.js';
 
-export const contextStorage = /* @__PURE__ */ new AsyncLocalStorage<{
+export type StepContext = {
   stepMetadata: StepMetadata;
   workflowMetadata: WorkflowMetadata;
   ops: Promise<void>[];
   closureVars?: Record<string, any>;
   encryptionKey?: CryptoKey;
-}>();
+};
+
+/**
+ * Process-wide singleton AsyncLocalStorage for step execution context.
+ *
+ * Uses `Symbol.for()` on globalThis to guarantee a single instance even when
+ * bundlers (e.g. Vercel's production bundler) create multiple copies of this
+ * module. Without this, `contextStorage.run()` in the step handler and
+ * `contextStorage.getStore()` in user code (via getWorkflowMetadata /
+ * getStepMetadata) can reference different AsyncLocalStorage instances,
+ * causing the store to appear empty.
+ * 
+ * Note that we were unable to reproduce this issue. This is a fix for the only synthetic way
+ * way in which we could get the builder to break with the reported error message, and 
+ * serves as defense-in-depth, since the change is otherwise safe.
+ *
+ * See: https://github.com/vercel/workflow/issues/1577
+ */
+const CONTEXT_STORAGE_SYMBOL = Symbol.for('WORKFLOW_STEP_CONTEXT_STORAGE');
+
+export const contextStorage: AsyncLocalStorage<StepContext> =
+  ((globalThis as any)[CONTEXT_STORAGE_SYMBOL] as
+    | AsyncLocalStorage<StepContext>
+    | undefined) ??
+  ((globalThis as any)[CONTEXT_STORAGE_SYMBOL] =
+    new AsyncLocalStorage<StepContext>());

workbench/example/workflows/99_e2e.ts

@@ -1502,3 +1502,51 @@ export async function importMetaUrlWorkflow() {
   'use workflow';
   return await checkImportMetaUrl();
 }
+
+//////////////////////////////////////////////////////////
+// Regression test for #1577:
+// getWorkflowMetadata()/getStepMetadata() called from a module-level helper
+// function (not directly inside the step body) must still have access to the
+// AsyncLocalStorage context.
+
+const withStrictMetadataCheck = async <T>(fn: () => Promise<T>) => {
+  const workflowMetadata = getWorkflowMetadata();
+  const stepMetadata = getStepMetadata();
+
+  return await fn().then((result) => ({
+    result,
+    workflowMetadata,
+    stepMetadata,
+  }));
+};
+
+async function metadataHelperStep(label: string): Promise<{
+  label: string;
+  workflowRunId: string;
+  stepId: string;
+  attempt: number;
+}> {
+  'use step';
+
+  const { workflowMetadata, stepMetadata } = await withStrictMetadataCheck(
+    async () => label
+  );
+
+  return {
+    label,
+    workflowRunId: workflowMetadata.workflowRunId,
+    stepId: stepMetadata.stepId,
+    attempt: stepMetadata.attempt,
+  };
+}
+
+export async function metadataFromHelperWorkflow(label: string): Promise<{
+  label: string;
+  workflowRunId: string;
+  stepId: string;
+  attempt: number;
+}> {
+  'use workflow';
+
+  return await metadataHelperStep(label);
+}