diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index 9d734d9a..61a2c47c 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -1,12 +1,14 @@ -import { createRequire } from 'node:module'; -import { join } from 'node:path'; - interface ContextManager { getStore(): T | undefined; run(value: T, fn: () => R): R; } +type AsyncLocalStorageLikeConstructor = new () => ContextManager; +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 +18,149 @@ 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 { + return ( + (typeof value === 'object' || typeof value === 'function') && + value !== null && + typeof (value as PromiseLike).then === 'function' + ); +} + +function createFallbackManager(): ContextManager { + let currentContext: any = undefined; + let activeAsyncRuns = 0; + + return { + getStore: () => currentContext, + run: (value: any, fn: () => R): R => { + if (activeAsyncRuns > 0) { + throw new Error(FALLBACK_CONCURRENCY_ERROR); + } + + const previousContext = currentContext; + currentContext = value; -function getNodeRequire(): NodeJS.Require { - if (typeof require === 'function') { - return require; + 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; + if (currentContext === value) { + 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; } - // We only require Node builtins, so any absolute path is valid as a createRequire base. - return createRequire(join(process.cwd(), '__axiom_require__.js')); + const mainModuleRequire = processRef?.mainModule?.require; + if (typeof mainModuleRequire === 'function') { + return mainModuleRequire; + } + + 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; + + try { + asyncHooksModule = legacyRequire('node:async_hooks'); + } catch { + try { + asyncHooksModule = legacyRequire('async_hooks'); + } catch { + return undefined; + } + } - let manager: ContextManager; + return getAsyncLocalStorageFromModule(asyncHooksModule); +} - if (isNodeJS) { +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 +169,13 @@ function getContextManager(): ContextManager { return manager; } -function createFallbackManager(): ContextManager { - let currentContext: any = null; - return { - getStore: () => currentContext, - run: (value: any, fn: () => R): R => { - const prev = currentContext; - currentContext = value; - try { - return fn(); - } finally { - currentContext = prev; - } - }, - }; -} - export function createAsyncHook(_name: string) { return { get(): T | undefined { - const manager = getContextManager(); - if (manager.getStore) { - return manager.getStore(); - } - return undefined; + return getContextManager().getStore(); }, run(value: T, fn: () => R): R { - const manager = getContextManager(); - return manager.run(value, fn); + return getContextManager().run(value, fn); }, }; } diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts new file mode 100644 index 00000000..cda8f1e0 --- /dev/null +++ b/packages/ai/test/evals/context/manager.test.ts @@ -0,0 +1,218 @@ +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') { + throw new Error('MODULE_NOT_FOUND'); + } + if (id === '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', + ); + }); + + it('handles thenables without a finally method in fallback mode', async () => { + processRef.getBuiltinModule = vi.fn(() => undefined); + (globalThis as any).require = undefined; + processRef.mainModule = { require: undefined }; + + vi.spyOn(console, 'warn').mockImplementation(() => {}); + const hook = createAsyncHook<{ requestId: string }>('test-context'); + + const thenable = { + then(onFulfilled: (value: string) => void) { + onFulfilled('ok'); + }, + }; + + const result = hook.run( + { requestId: 'req-thenable' }, + () => thenable as any, + ) as unknown as Promise; + + await expect(result).resolves.toBe('ok'); + expect(hook.get()).toBeUndefined(); + }); + + it('rejects nested runs while a fallback run is active', async () => { + processRef.getBuiltinModule = vi.fn(() => undefined); + (globalThis as any).require = undefined; + processRef.mainModule = { require: undefined }; + + vi.spyOn(console, 'warn').mockImplementation(() => {}); + const hook = createAsyncHook<{ requestId: string }>('test-context'); + + let releaseFirst: (() => void) | undefined; + const firstRun = hook.run({ requestId: 'first' }, async () => { + await new Promise((resolve) => { + releaseFirst = resolve; + }); + expect(hook.get()).toEqual({ requestId: 'first' }); + }); + + await Promise.resolve(); + + const secondRun = vi.fn(async () => { + await Promise.resolve(); + }); + + expect(() => { + hook.run({ requestId: 'second' }, secondRun); + }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); + + expect(secondRun).not.toHaveBeenCalled(); + + releaseFirst?.(); + await firstRun; + }); + + it('blocks nested async runs before they can leak context', async () => { + processRef.getBuiltinModule = vi.fn(() => undefined); + (globalThis as any).require = undefined; + processRef.mainModule = { require: undefined }; + + vi.spyOn(console, 'warn').mockImplementation(() => {}); + const hook = createAsyncHook<{ requestId: string }>('test-context'); + + expect(() => { + hook.run({ requestId: 'outer' }, () => { + return hook.run({ requestId: 'inner' }, async () => { + await Promise.resolve(); + return 'ok'; + }); + }); + }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); + + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(hook.get()).toBeUndefined(); + }); + + it('allows nested synchronous runs when no async fallback run is active', () => { + processRef.getBuiltinModule = vi.fn(() => undefined); + (globalThis as any).require = undefined; + processRef.mainModule = { require: undefined }; + + vi.spyOn(console, 'warn').mockImplementation(() => {}); + const hook = createAsyncHook<{ requestId: string }>('test-context'); + + hook.run({ requestId: 'outer' }, () => { + const nested = hook.run({ requestId: 'inner' }, () => hook.get()); + expect(nested).toEqual({ requestId: 'inner' }); + expect(hook.get()).toEqual({ requestId: 'outer' }); + }); + + expect(hook.get()).toBeUndefined(); + }); + + it('detects and prevents context leakage in fallback mode', async () => { + processRef.getBuiltinModule = vi.fn(() => undefined); + (globalThis as any).require = undefined; + processRef.mainModule = { require: undefined }; + + vi.spyOn(console, 'warn').mockImplementation(() => {}); + const hook = createAsyncHook<{ requestId: string }>('test-context'); + + let resolveInner: (v: any) => void = () => {}; + const innerPromise = new Promise((r) => { + resolveInner = r; + }); + + hook.run({ requestId: 'outer' }, () => { + // Start an async run but don't return it + hook.run({ requestId: 'inner' }, () => innerPromise); + }); + + // Outer run finished synchronously. Context should be restored to undefined. + expect(hook.get()).toBeUndefined(); + + resolveInner(null); + await innerPromise; + await new Promise((resolve) => setTimeout(resolve, 0)); // Ensure finally runs + + // After inner promise finished, context should still be undefined, not 'outer'. + expect(hook.get()).toBeUndefined(); + }); +}); diff --git a/packages/ai/tsup.config.ts b/packages/ai/tsup.config.ts index 1a220301..b0390f47 100644 --- a/packages/ai/tsup.config.ts +++ b/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([