fix(ai): replace static async_hooks require in eval context

packages/ai/src/evals/context/manager.ts

@@ -1,13 +1,17 @@
-import { createRequire } from 'node:module';
-import { join } from 'node:path';
-
 interface ContextManager<T = any> {
   getStore(): T | undefined;
   run<R>(value: T, fn: () => R): R;
 }
 
+type AsyncLocalStorageLikeConstructor = new <T = any>() => ContextManager<T>;
+
+type NodeRequireLike = (id: string) => unknown;
+
 const CONTEXT_MANAGER_SYMBOL = Symbol.for('axiom.context_manager');
 
+const FALLBACK_CONCURRENCY_ERROR =
+  'AsyncLocalStorage fallback does not support concurrent async contexts';
+
 function getGlobalContextManager(): ContextManager | undefined {
   return (globalThis as any)[CONTEXT_MANAGER_SYMBOL];
 }
@@ -16,47 +20,141 @@ function setGlobalContextManager(manager: ContextManager): void {
   (globalThis as any)[CONTEXT_MANAGER_SYMBOL] = manager;
 }
 
-const isNodeJS = typeof process !== 'undefined' && !!process.versions?.node;
+function isPromiseLike(value: unknown): value is PromiseLike<unknown> {
+  return (
+    (typeof value === 'object' || typeof value === 'function') &&
+    value !== null &&
+    typeof (value as PromiseLike<unknown>).then === 'function'
+  );
+}
+
+function createFallbackManager(): ContextManager {
+  let currentContext: any = undefined;
+  let activeAsyncRuns = 0;
+
+  return {
+    getStore: () => currentContext,
+    run: <R>(value: any, fn: () => R): R => {
+      const previousContext = currentContext;
+      currentContext = value;
+
+      let result: R;
+      try {
+        result = fn();
+      } catch (error) {
+        currentContext = previousContext;
+        throw error;
+      }
+
+      if (isPromiseLike(result)) {
+        if (activeAsyncRuns > 0) {
+          currentContext = previousContext;
+          throw new Error(FALLBACK_CONCURRENCY_ERROR);
+        }
+
+        activeAsyncRuns += 1;
+
+        return Promise.resolve(result).finally(() => {
+          activeAsyncRuns -= 1;
+          currentContext = previousContext;
+        }) as R;
+      }
+
+      currentContext = previousContext;
+      return result;
+    },
+  };
+}
+
+function getAsyncLocalStorageFromModule(module: unknown): AsyncLocalStorageLikeConstructor | undefined {
+  const asyncLocalStorageCtor = (module as { AsyncLocalStorage?: AsyncLocalStorageLikeConstructor })
+    ?.AsyncLocalStorage;
+  return typeof asyncLocalStorageCtor === 'function' ? asyncLocalStorageCtor : undefined;
+}
+
+function getLegacyNodeRequire(): NodeRequireLike | undefined {
+  const processRef = (globalThis as any).process as { mainModule?: { require?: NodeRequireLike } };
+  const globalRequire = (globalThis as any).require as NodeRequireLike | undefined;
+
+  if (typeof globalRequire === 'function') {
+    return globalRequire;
+  }
 
-function getNodeRequire(): NodeJS.Require {
-  if (typeof require === 'function') {
-    return require;
+  const mainModuleRequire = processRef?.mainModule?.require;
+  if (typeof mainModuleRequire === 'function') {
+    return mainModuleRequire;
   }
 
-  // We only require Node builtins, so any absolute path is valid as a createRequire base.
-  return createRequire(join(process.cwd(), '__axiom_require__.js'));
+  return undefined;
 }
 
-function getContextManager(): ContextManager {
-  // Check global Symbol registry cache first (shared across VM contexts)
-  const existing = getGlobalContextManager();
-  if (existing) return existing;
+function getAsyncLocalStorageFromLegacyRequire(
+  legacyRequire: NodeRequireLike,
+): AsyncLocalStorageLikeConstructor | undefined {
+  let asyncHooksModule: unknown;
 
-  let manager: ContextManager;
+  try {
+    asyncHooksModule = legacyRequire('node:async_hooks');
+  } catch {
+    try {
+      asyncHooksModule = legacyRequire('async_hooks');
+    } catch {
+      return undefined;
+    }
+  }
 
-  if (isNodeJS) {
+  return getAsyncLocalStorageFromModule(asyncHooksModule);
+}
+
+function getAsyncLocalStorageConstructor(): AsyncLocalStorageLikeConstructor | undefined {
+  const getBuiltinModule = (globalThis as any).process?.getBuiltinModule as
+    | ((id: string) => unknown)
+    | undefined;
+
+  if (typeof getBuiltinModule === 'function') {
     try {
-      // Resolve AsyncLocalStorage in both ESM and CJS Node contexts without bundler interference
-      let AsyncLocalStorage: any;
+      const asyncHooksModule =
+        getBuiltinModule('node:async_hooks') ?? getBuiltinModule('async_hooks');
 
-      // Obtain require in both CJS and ESM Node runtimes without relying on tsup shims.
-      const req = getNodeRequire();
-      try {
-        AsyncLocalStorage = req('node:async_hooks').AsyncLocalStorage;
-      } catch {
-        AsyncLocalStorage = req('async_hooks').AsyncLocalStorage;
+      const asyncLocalStorageCtor = getAsyncLocalStorageFromModule(asyncHooksModule);
+      if (asyncLocalStorageCtor) {
+        return asyncLocalStorageCtor;
       }
+    } catch (error) {
+      console.warn('Failed to load AsyncLocalStorage from node:async_hooks:', error);
+    }
+  }
 
-      manager = new AsyncLocalStorage();
+  const legacyRequire = getLegacyNodeRequire();
+  if (typeof legacyRequire === 'function') {
+    try {
+      const asyncLocalStorageCtor = getAsyncLocalStorageFromLegacyRequire(legacyRequire);
+      if (asyncLocalStorageCtor) {
+        return asyncLocalStorageCtor;
+      }
     } catch (error) {
-      // Fallback if AsyncLocalStorage cannot be loaded
-      console.warn('AsyncLocalStorage not available, using fallback context manager:', error);
-      manager = createFallbackManager();
+      console.warn('Failed to load AsyncLocalStorage via legacy require fallback:', error);
     }
-  } else {
-    // Browser/CF Workers - simple fallback (no warning needed here)
+  }
+
+  const globalAsyncLocalStorage = (globalThis as any).AsyncLocalStorage;
+  if (typeof globalAsyncLocalStorage === 'function') {
+    return globalAsyncLocalStorage as AsyncLocalStorageLikeConstructor;
+  }
+
+  return undefined;
+}
+
+function getContextManager(): ContextManager {
+  // Check global Symbol registry cache first (shared across VM contexts)
+  const existing = getGlobalContextManager();
+  if (existing) return existing;
+
+  const AsyncLocalStorageCtor = getAsyncLocalStorageConstructor();
+  const manager = AsyncLocalStorageCtor ? new AsyncLocalStorageCtor() : createFallbackManager();
+
+  if (!AsyncLocalStorageCtor) {
     console.warn('AsyncLocalStorage not available, using fallback context manager');
-    manager = createFallbackManager();
   }
 
   // Cache using Symbol to share across VM contexts
@@ -65,34 +163,13 @@ function getContextManager(): ContextManager {
   return manager;
 }
 
-function createFallbackManager(): ContextManager {
-  let currentContext: any = null;
-  return {
-    getStore: () => currentContext,
-    run: <R>(value: any, fn: () => R): R => {
-      const prev = currentContext;
-      currentContext = value;
-      try {
-        return fn();
-      } finally {
-        currentContext = prev;
-      }
-    },
-  };
-}
-
 export function createAsyncHook<T>(_name: string) {
   return {
     get(): T | undefined {
-      const manager = getContextManager();
-      if (manager.getStore) {
-        return manager.getStore();
-      }
-      return undefined;
+      return getContextManager().getStore();
     },
     run<R>(value: T, fn: () => R): R {
-      const manager = getContextManager();
-      return manager.run(value, fn);
+      return getContextManager().run(value, fn);
     },
   };
 }
@@ -105,3 +182,4 @@ export function createAsyncHook<T>(_name: string) {
 export function __resetContextManagerForTests(): void {
   delete (globalThis as any)[CONTEXT_MANAGER_SYMBOL];
 }
+

packages/ai/test/evals/context/manager.test.ts

@@ -0,0 +1,90 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { __resetContextManagerForTests, createAsyncHook } from '../../../src/evals/context/manager';
+
+type BuiltinModuleLoader = ((id: string) => unknown) | undefined;
+type NodeRequireLike = ((id: string) => unknown) | undefined;
+
+type ProcessRef = {
+  getBuiltinModule?: BuiltinModuleLoader;
+  mainModule?: {
+    require?: NodeRequireLike;
+  };
+};
+
+const processRef = (globalThis as any).process as ProcessRef;
+
+const originalGetBuiltinModule = processRef.getBuiltinModule;
+const originalMainModule = processRef.mainModule;
+const originalGlobalRequire = (globalThis as any).require as NodeRequireLike;
+
+describe('eval context manager', () => {
+  beforeEach(() => {
+    __resetContextManagerForTests();
+  });
+
+  afterEach(() => {
+    processRef.getBuiltinModule = originalGetBuiltinModule;
+    processRef.mainModule = originalMainModule;
+    (globalThis as any).require = originalGlobalRequire;
+    __resetContextManagerForTests();
+    vi.restoreAllMocks();
+  });
+
+  it('uses AsyncLocalStorage when getBuiltinModule is available', async () => {
+    processRef.getBuiltinModule = vi.fn((id: string) => {
+      if (id === 'node:async_hooks') {
+        return { AsyncLocalStorage };
+      }
+      return undefined;
+    });
+
+    const hook = createAsyncHook<{ requestId: string }>('test-context');
+
+    await hook.run({ requestId: 'req-123' }, async () => {
+      await Promise.resolve();
+      expect(hook.get()).toEqual({ requestId: 'req-123' });
+    });
+
+    expect(hook.get()).toBeUndefined();
+  });
+
+  it('uses legacy require fallback when getBuiltinModule is unavailable', async () => {
+    processRef.getBuiltinModule = undefined;
+    (globalThis as any).require = vi.fn((id: string) => {
+      if (id === 'node:async_hooks') {
+        return { AsyncLocalStorage };
+      }
+      return undefined;
+    });
+
+    const fallbackWarn = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    const hook = createAsyncHook<{ requestId: string }>('test-context');
+
+    await hook.run({ requestId: 'req-legacy' }, async () => {
+      await Promise.resolve();
+      expect(hook.get()).toEqual({ requestId: 'req-legacy' });
+    });
+
+    expect(fallbackWarn).not.toHaveBeenCalledWith(
+      'AsyncLocalStorage not available, using fallback context manager',
+    );
+  });
+
+  it('falls back when AsyncLocalStorage is unavailable', async () => {
+    processRef.getBuiltinModule = vi.fn(() => undefined);
+    (globalThis as any).require = undefined;
+    processRef.mainModule = { require: undefined };
+
+    const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    const hook = createAsyncHook<{ requestId: string }>('test-context');
+
+    await hook.run({ requestId: 'req-456' }, async () => {
+      await Promise.resolve();
+      expect(hook.get()).toEqual({ requestId: 'req-456' });
+    });
+
+    expect(hook.get()).toBeUndefined();
+    expect(warn).toHaveBeenCalledWith('AsyncLocalStorage not available, using fallback context manager');
+  });
+});

packages/ai/tsup.config.ts

@@ -13,11 +13,8 @@ const sharedConfig = {
     'c12',
     'defu',
     'vite-tsconfig-paths',
-    // Ensure Node builtins used via createRequire stay external in ESM bundle
-    'async_hooks',
+    // Keep async_hooks external so the ESM build preserves Node's builtin import.
     'node:async_hooks',
-    'module',
-    'node:module',
   ],
   dts: true,
   sourcemap: true,
@@ -27,6 +24,7 @@ const sharedConfig = {
   define: {
     __SDK_VERSION__: JSON.stringify(pkg.version),
   },
+  removeNodeProtocol: false,
 };
 
 export default defineConfig([