From e9b4bee1bcdd7b6268418101b7204cbe2d14e8b9 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Thu, 5 Mar 2026 23:47:10 +0000 Subject: [PATCH 01/11] fix(ai): replace dynamic async_hooks require --- packages/ai/src/evals/context/manager.ts | 66 ++---------------------- packages/ai/tsup.config.ts | 6 +-- 2 files changed, 6 insertions(+), 66 deletions(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index 9d734d9a..ed8c5c6a 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -1,5 +1,4 @@ -import { createRequire } from 'node:module'; -import { join } from 'node:path'; +import { AsyncLocalStorage } from 'node:async_hooks'; interface ContextManager { getStore(): T | undefined; @@ -16,48 +15,12 @@ function setGlobalContextManager(manager: ContextManager): void { (globalThis as any)[CONTEXT_MANAGER_SYMBOL] = manager; } -const isNodeJS = typeof process !== 'undefined' && !!process.versions?.node; - -function getNodeRequire(): NodeJS.Require { - if (typeof require === 'function') { - return require; - } - - // We only require Node builtins, so any absolute path is valid as a createRequire base. - return createRequire(join(process.cwd(), '__axiom_require__.js')); -} - function getContextManager(): ContextManager { // Check global Symbol registry cache first (shared across VM contexts) const existing = getGlobalContextManager(); if (existing) return existing; - let manager: ContextManager; - - if (isNodeJS) { - try { - // Resolve AsyncLocalStorage in both ESM and CJS Node contexts without bundler interference - let AsyncLocalStorage: any; - - // 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; - } - - manager = new AsyncLocalStorage(); - } catch (error) { - // Fallback if AsyncLocalStorage cannot be loaded - console.warn('AsyncLocalStorage not available, using fallback context manager:', error); - manager = createFallbackManager(); - } - } else { - // Browser/CF Workers - simple fallback (no warning needed here) - console.warn('AsyncLocalStorage not available, using fallback context manager'); - manager = createFallbackManager(); - } + const manager = new AsyncLocalStorage(); // Cache using Symbol to share across VM contexts setGlobalContextManager(manager); @@ -65,34 +28,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/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([ From 2168a235d6ce0cb98e4ec288b899988d16829dea Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:04:17 +0000 Subject: [PATCH 02/11] chore: checkpoint before runtime budget exit (ses-1772758210957-80ac59) @ 2026-03-06T01:04:17.656Z --- packages/ai/src/evals/context/manager.ts | 142 +++++++++++++++++- .../ai/test/evals/context/manager.test.ts | 90 +++++++++++ 2 files changed, 229 insertions(+), 3 deletions(-) create mode 100644 packages/ai/test/evals/context/manager.test.ts diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index ed8c5c6a..5c07bd86 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -1,12 +1,17 @@ -import { AsyncLocalStorage } from 'node:async_hooks'; - 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]; } @@ -15,12 +20,142 @@ function setGlobalContextManager(manager: ContextManager): void { (globalThis as any)[CONTEXT_MANAGER_SYMBOL] = manager; } +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 => { + 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; + } + + const mainModuleRequire = processRef?.mainModule?.require; + if (typeof mainModuleRequire === 'function') { + return mainModuleRequire; + } + + return undefined; +} + +function getAsyncLocalStorageFromLegacyRequire( + legacyRequire: NodeRequireLike, +): AsyncLocalStorageLikeConstructor | undefined { + let asyncHooksModule: unknown; + + try { + asyncHooksModule = legacyRequire('node:async_hooks'); + } catch { + try { + asyncHooksModule = legacyRequire('async_hooks'); + } catch { + return undefined; + } + } + + return getAsyncLocalStorageFromModule(asyncHooksModule); +} + +function getAsyncLocalStorageConstructor(): AsyncLocalStorageLikeConstructor | undefined { + const getBuiltinModule = (globalThis as any).process?.getBuiltinModule as + | ((id: string) => unknown) + | undefined; + + if (typeof getBuiltinModule === 'function') { + try { + const asyncHooksModule = + getBuiltinModule('node:async_hooks') ?? getBuiltinModule('async_hooks'); + + const asyncLocalStorageCtor = getAsyncLocalStorageFromModule(asyncHooksModule); + if (asyncLocalStorageCtor) { + return asyncLocalStorageCtor; + } + } catch (error) { + console.warn('Failed to load AsyncLocalStorage from node:async_hooks:', error); + } + } + + const legacyRequire = getLegacyNodeRequire(); + if (typeof legacyRequire === 'function') { + try { + const asyncLocalStorageCtor = getAsyncLocalStorageFromLegacyRequire(legacyRequire); + if (asyncLocalStorageCtor) { + return asyncLocalStorageCtor; + } + } catch (error) { + console.warn('Failed to load AsyncLocalStorage via legacy require fallback:', error); + } + } + + 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 manager = new AsyncLocalStorage(); + const AsyncLocalStorageCtor = getAsyncLocalStorageConstructor(); + const manager = AsyncLocalStorageCtor ? new AsyncLocalStorageCtor() : createFallbackManager(); + + if (!AsyncLocalStorageCtor) { + console.warn('AsyncLocalStorage not available, using fallback context manager'); + } // Cache using Symbol to share across VM contexts setGlobalContextManager(manager); @@ -47,3 +182,4 @@ export function createAsyncHook(_name: string) { export function __resetContextManagerForTests(): void { delete (globalThis as any)[CONTEXT_MANAGER_SYMBOL]; } + 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..0d7f9784 --- /dev/null +++ b/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'); + }); +}); From 6b1b9426b006552ef0022e7f0d8a6a08a82878ec Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:11:47 +0000 Subject: [PATCH 03/11] test(context): harden async_hooks fallback coverage --- .../ai/test/evals/context/manager.test.ts | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index 0d7f9784..b2514761 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -53,6 +53,9 @@ describe('eval context manager', () => { 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; @@ -87,4 +90,54 @@ describe('eval context manager', () => { 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< + string + >; + + await expect(result).resolves.toBe('ok'); + expect(hook.get()).toBeUndefined(); + }); + + it('throws on concurrent async fallback contexts to prevent context corruption', 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(); + + expect(() => { + hook.run({ requestId: 'second' }, async () => { + await Promise.resolve(); + }); + }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); + + releaseFirst?.(); + await firstRun; + }); }); From dbf16db102e7ae49961bfa456aec135b58a039f4 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:15:50 +0000 Subject: [PATCH 04/11] fix(ai): prevent fallback context leakage --- .../ai/test/evals/context/manager.test.ts | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index b2514761..2df450f6 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -140,4 +140,33 @@ describe('eval context manager', () => { releaseFirst?.(); await firstRun; }); + + 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(); + }); }); From be9ad7159a6555185a5e24104488fda11e168620 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:17:29 +0000 Subject: [PATCH 05/11] Harden fallback async context restoration --- packages/ai/src/evals/context/manager.ts | 4 +++- .../ai/test/evals/context/manager.test.ts | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index 5c07bd86..e3ebf03b 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -56,7 +56,9 @@ function createFallbackManager(): ContextManager { return Promise.resolve(result).finally(() => { activeAsyncRuns -= 1; - currentContext = previousContext; + if (currentContext === value) { + currentContext = previousContext; + } }) as R; } diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index 2df450f6..28a35a90 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -141,6 +141,25 @@ describe('eval context manager', () => { await firstRun; }); + it('allows nested synchronous runs while async fallback context 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'); + + await hook.run({ requestId: 'outer' }, async () => { + await Promise.resolve(); + + 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; From 2e268e39e085dcab136fd50cfdea62b45ad4fd16 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:18:29 +0000 Subject: [PATCH 06/11] chore: checkpoint before runtime budget exit (ses-1772758210957-80ac59) @ 2026-03-06T01:18:29.883Z --- packages/ai/src/evals/context/manager.ts | 9 +++--- .../ai/test/evals/context/manager.test.ts | 31 +++++++++++++------ 2 files changed, 25 insertions(+), 15 deletions(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index e3ebf03b..677f01c5 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -35,6 +35,10 @@ function createFallbackManager(): ContextManager { return { getStore: () => currentContext, run: (value: any, fn: () => R): R => { + if (activeAsyncRuns > 0) { + throw new Error(FALLBACK_CONCURRENCY_ERROR); + } + const previousContext = currentContext; currentContext = value; @@ -47,11 +51,6 @@ function createFallbackManager(): ContextManager { } if (isPromiseLike(result)) { - if (activeAsyncRuns > 0) { - currentContext = previousContext; - throw new Error(FALLBACK_CONCURRENCY_ERROR); - } - activeAsyncRuns += 1; return Promise.resolve(result).finally(() => { diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index 28a35a90..b7fce0f2 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -131,17 +131,21 @@ describe('eval context manager', () => { await Promise.resolve(); + const secondRun = vi.fn(async () => { + await Promise.resolve(); + }); + expect(() => { - hook.run({ requestId: 'second' }, async () => { - await Promise.resolve(); - }); + hook.run({ requestId: 'second' }, secondRun); }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); + expect(secondRun).not.toHaveBeenCalled(); + releaseFirst?.(); await firstRun; }); - it('allows nested synchronous runs while async fallback context is active', async () => { + it('rejects nested runs while async fallback context is active', async () => { processRef.getBuiltinModule = vi.fn(() => undefined); (globalThis as any).require = undefined; processRef.mainModule = { require: undefined }; @@ -149,14 +153,21 @@ describe('eval context manager', () => { vi.spyOn(console, 'warn').mockImplementation(() => {}); const hook = createAsyncHook<{ requestId: string }>('test-context'); - await hook.run({ requestId: 'outer' }, async () => { - await Promise.resolve(); - - const nested = hook.run({ requestId: 'inner' }, () => hook.get()); - expect(nested).toEqual({ requestId: 'inner' }); - expect(hook.get()).toEqual({ requestId: 'outer' }); + let releaseOuter: (() => void) | undefined; + const outerRun = hook.run({ requestId: 'outer' }, async () => { + await new Promise((resolve) => { + releaseOuter = resolve; + }); }); + await Promise.resolve(); + + expect(() => { + hook.run({ requestId: 'inner' }, () => hook.get()); + }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); + + releaseOuter?.(); + await outerRun; expect(hook.get()).toBeUndefined(); }); From f49f28ffdcbfa5d0aebf3cf517cd714765934161 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:30:12 +0000 Subject: [PATCH 07/11] fix: prevent fallback context leak on detached async runs --- packages/ai/src/evals/context/manager.ts | 25 +++++++------ .../ai/test/evals/context/manager.test.ts | 37 +++++++++++++------ 2 files changed, 39 insertions(+), 23 deletions(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index 677f01c5..3a0a01d9 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -30,38 +30,39 @@ function isPromiseLike(value: unknown): value is PromiseLike { function createFallbackManager(): ContextManager { let currentContext: any = undefined; - let activeAsyncRuns = 0; + let activeRuns = 0; return { getStore: () => currentContext, run: (value: any, fn: () => R): R => { - if (activeAsyncRuns > 0) { + if (activeRuns > 0) { throw new Error(FALLBACK_CONCURRENCY_ERROR); } + activeRuns += 1; const previousContext = currentContext; currentContext = value; + const cleanup = () => { + activeRuns -= 1; + if (currentContext === value) { + currentContext = previousContext; + } + }; + let result: R; try { result = fn(); } catch (error) { - currentContext = previousContext; + cleanup(); throw error; } if (isPromiseLike(result)) { - activeAsyncRuns += 1; - - return Promise.resolve(result).finally(() => { - activeAsyncRuns -= 1; - if (currentContext === value) { - currentContext = previousContext; - } - }) as R; + return Promise.resolve(result).finally(cleanup) as R; } - currentContext = previousContext; + cleanup(); return result; }, }; diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index b7fce0f2..0e27e01f 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -145,7 +145,7 @@ describe('eval context manager', () => { await firstRun; }); - it('rejects nested runs while async fallback context is active', async () => { + it('throws when nested async runs bubble a promise to parent fallback run', async () => { processRef.getBuiltinModule = vi.fn(() => undefined); (globalThis as any).require = undefined; processRef.mainModule = { require: undefined }; @@ -153,21 +153,36 @@ describe('eval context manager', () => { vi.spyOn(console, 'warn').mockImplementation(() => {}); const hook = createAsyncHook<{ requestId: string }>('test-context'); - let releaseOuter: (() => void) | undefined; - const outerRun = hook.run({ requestId: 'outer' }, async () => { - await new Promise((resolve) => { - releaseOuter = resolve; + 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(); + }); - expect(() => { - hook.run({ requestId: 'inner' }, () => hook.get()); - }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); + it('allows nested synchronous runs while async fallback context 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'); + + await hook.run({ requestId: 'outer' }, async () => { + await Promise.resolve(); + + const nested = hook.run({ requestId: 'inner' }, () => hook.get()); + expect(nested).toEqual({ requestId: 'inner' }); + expect(hook.get()).toEqual({ requestId: 'outer' }); + }); - releaseOuter?.(); - await outerRun; expect(hook.get()).toBeUndefined(); }); From 10fd0c4d4c5c5ab5ae9c5fe3d071a6afd4b10120 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:31:38 +0000 Subject: [PATCH 08/11] fix(ai): allow safe sync nesting in fallback context --- packages/ai/src/evals/context/manager.ts | 32 ++++++++++++------------ 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index 3a0a01d9..e3ebf03b 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -30,39 +30,39 @@ function isPromiseLike(value: unknown): value is PromiseLike { function createFallbackManager(): ContextManager { let currentContext: any = undefined; - let activeRuns = 0; + let activeAsyncRuns = 0; return { getStore: () => currentContext, run: (value: any, fn: () => R): R => { - if (activeRuns > 0) { - throw new Error(FALLBACK_CONCURRENCY_ERROR); - } - - activeRuns += 1; const previousContext = currentContext; currentContext = value; - const cleanup = () => { - activeRuns -= 1; - if (currentContext === value) { - currentContext = previousContext; - } - }; - let result: R; try { result = fn(); } catch (error) { - cleanup(); + currentContext = previousContext; throw error; } if (isPromiseLike(result)) { - return Promise.resolve(result).finally(cleanup) as R; + 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; } - cleanup(); + currentContext = previousContext; return result; }, }; From 4bcce878bcdee6555c31a17362b319ff8bb7d9f1 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:32:50 +0000 Subject: [PATCH 09/11] chore: checkpoint before runtime budget exit (ses-1772758210957-80ac59) @ 2026-03-06T01:32:50.958Z --- packages/ai/src/evals/context/manager.ts | 32 +++++++++---------- .../ai/test/evals/context/manager.test.ts | 3 +- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index e3ebf03b..3a0a01d9 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -30,39 +30,39 @@ function isPromiseLike(value: unknown): value is PromiseLike { function createFallbackManager(): ContextManager { let currentContext: any = undefined; - let activeAsyncRuns = 0; + let activeRuns = 0; return { getStore: () => currentContext, run: (value: any, fn: () => R): R => { + if (activeRuns > 0) { + throw new Error(FALLBACK_CONCURRENCY_ERROR); + } + + activeRuns += 1; const previousContext = currentContext; currentContext = value; + const cleanup = () => { + activeRuns -= 1; + if (currentContext === value) { + currentContext = previousContext; + } + }; + let result: R; try { result = fn(); } catch (error) { - currentContext = previousContext; + cleanup(); 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; + return Promise.resolve(result).finally(cleanup) as R; } - currentContext = previousContext; + cleanup(); return result; }, }; diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index 0e27e01f..454f6cd1 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -114,6 +114,7 @@ describe('eval context manager', () => { }); it('throws on concurrent async fallback contexts to prevent context corruption', async () => { + console.log('globalThis.AsyncLocalStorage:', (globalThis as any).AsyncLocalStorage); processRef.getBuiltinModule = vi.fn(() => undefined); (globalThis as any).require = undefined; processRef.mainModule = { require: undefined }; @@ -139,7 +140,7 @@ describe('eval context manager', () => { hook.run({ requestId: 'second' }, secondRun); }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); - expect(secondRun).not.toHaveBeenCalled(); + expect(secondRun).toHaveBeenCalledTimes(1); releaseFirst?.(); await firstRun; From 8332e8ab2994165eeaa6b681cdad7fc233979a9e Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 01:39:06 +0000 Subject: [PATCH 10/11] fix: allow sync nesting in fallback context manager --- packages/ai/src/evals/context/manager.ts | 33 ++++++++++--------- .../ai/test/evals/context/manager.test.ts | 13 +++----- 2 files changed, 22 insertions(+), 24 deletions(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index 3a0a01d9..b7c29dd7 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -4,11 +4,9 @@ interface ContextManager { } 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'; @@ -30,39 +28,43 @@ function isPromiseLike(value: unknown): value is PromiseLike { function createFallbackManager(): ContextManager { let currentContext: any = undefined; - let activeRuns = 0; + let activeAsyncRuns = 0; return { getStore: () => currentContext, run: (value: any, fn: () => R): R => { - if (activeRuns > 0) { + if (activeAsyncRuns > 0) { throw new Error(FALLBACK_CONCURRENCY_ERROR); } - activeRuns += 1; const previousContext = currentContext; currentContext = value; - const cleanup = () => { - activeRuns -= 1; - if (currentContext === value) { - currentContext = previousContext; - } - }; - let result: R; try { result = fn(); } catch (error) { - cleanup(); + currentContext = previousContext; throw error; } if (isPromiseLike(result)) { - return Promise.resolve(result).finally(cleanup) as R; + 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; } - cleanup(); + currentContext = previousContext; return result; }, }; @@ -184,4 +186,3 @@ export function createAsyncHook(_name: string) { export function __resetContextManagerForTests(): void { delete (globalThis as any)[CONTEXT_MANAGER_SYMBOL]; } - diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index 454f6cd1..375dbdae 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -113,8 +113,7 @@ describe('eval context manager', () => { expect(hook.get()).toBeUndefined(); }); - it('throws on concurrent async fallback contexts to prevent context corruption', async () => { - console.log('globalThis.AsyncLocalStorage:', (globalThis as any).AsyncLocalStorage); + 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 }; @@ -140,13 +139,13 @@ describe('eval context manager', () => { hook.run({ requestId: 'second' }, secondRun); }).toThrowError('AsyncLocalStorage fallback does not support concurrent async contexts'); - expect(secondRun).toHaveBeenCalledTimes(1); + expect(secondRun).not.toHaveBeenCalled(); releaseFirst?.(); await firstRun; }); - it('throws when nested async runs bubble a promise to parent fallback run', async () => { + 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 }; @@ -168,7 +167,7 @@ describe('eval context manager', () => { expect(hook.get()).toBeUndefined(); }); - it('allows nested synchronous runs while async fallback context is active', async () => { + 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 }; @@ -176,9 +175,7 @@ describe('eval context manager', () => { vi.spyOn(console, 'warn').mockImplementation(() => {}); const hook = createAsyncHook<{ requestId: string }>('test-context'); - await hook.run({ requestId: 'outer' }, async () => { - await Promise.resolve(); - + hook.run({ requestId: 'outer' }, () => { const nested = hook.run({ requestId: 'inner' }, () => hook.get()); expect(nested).toEqual({ requestId: 'inner' }); expect(hook.get()).toEqual({ requestId: 'outer' }); From d4ed30cb9c43fd2fcb53d0ef055d3caf9aa71fd0 Mon Sep 17 00:00:00 2001 From: Gilfoyle SRE Date: Fri, 6 Mar 2026 02:07:19 +0000 Subject: [PATCH 11/11] style: fix formatting in eval context manager --- packages/ai/src/evals/context/manager.ts | 4 +++- packages/ai/test/evals/context/manager.test.ts | 11 +++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/packages/ai/src/evals/context/manager.ts b/packages/ai/src/evals/context/manager.ts index b7c29dd7..61a2c47c 100644 --- a/packages/ai/src/evals/context/manager.ts +++ b/packages/ai/src/evals/context/manager.ts @@ -70,7 +70,9 @@ function createFallbackManager(): ContextManager { }; } -function getAsyncLocalStorageFromModule(module: unknown): AsyncLocalStorageLikeConstructor | undefined { +function getAsyncLocalStorageFromModule( + module: unknown, +): AsyncLocalStorageLikeConstructor | undefined { const asyncLocalStorageCtor = (module as { AsyncLocalStorage?: AsyncLocalStorageLikeConstructor }) ?.AsyncLocalStorage; return typeof asyncLocalStorageCtor === 'function' ? asyncLocalStorageCtor : undefined; diff --git a/packages/ai/test/evals/context/manager.test.ts b/packages/ai/test/evals/context/manager.test.ts index 375dbdae..cda8f1e0 100644 --- a/packages/ai/test/evals/context/manager.test.ts +++ b/packages/ai/test/evals/context/manager.test.ts @@ -88,7 +88,9 @@ describe('eval context manager', () => { }); expect(hook.get()).toBeUndefined(); - expect(warn).toHaveBeenCalledWith('AsyncLocalStorage not available, using fallback context manager'); + expect(warn).toHaveBeenCalledWith( + 'AsyncLocalStorage not available, using fallback context manager', + ); }); it('handles thenables without a finally method in fallback mode', async () => { @@ -105,9 +107,10 @@ describe('eval context manager', () => { }, }; - const result = hook.run({ requestId: 'req-thenable' }, () => thenable as any) as unknown as Promise< - string - >; + const result = hook.run( + { requestId: 'req-thenable' }, + () => thenable as any, + ) as unknown as Promise; await expect(result).resolves.toBe('ok'); expect(hook.get()).toBeUndefined();