diff --git a/.changeset/large-gifts-go.md b/.changeset/large-gifts-go.md new file mode 100644 index 000000000..00b3dd5ff --- /dev/null +++ b/.changeset/large-gifts-go.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +ref: Fork import-in-the-middle and require-in-the-middle diff --git a/js/src/auto-instrumentations/bundler/plugin.ts b/js/src/auto-instrumentations/bundler/plugin.ts index dda1b7490..c1d31d451 100644 --- a/js/src/auto-instrumentations/bundler/plugin.ts +++ b/js/src/auto-instrumentations/bundler/plugin.ts @@ -1,11 +1,20 @@ import { createUnplugin } from "unplugin"; import { create, type InstrumentationConfig } from "../orchestrion-js"; -import { extname, join, sep } from "path"; +import { dirname, extname, join, resolve, sep } from "path"; import { readFileSync } from "fs"; import { fileURLToPath } from "url"; import moduleDetailsFromPath from "module-details-from-path"; import { getDefaultInstrumentationConfigs } from "../configs/all"; import { applySpecialCasePatch } from "../loader/special-case-patches"; +import { + buildTopLevelImportHookSourceWrapper, + getDefaultTopLevelImportHooks, + type TopLevelImportHookTarget, +} from "../loader/top-level-export-patches"; +import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config"; + +const TOP_LEVEL_ORIGINAL_IMPORT_PREFIX = "braintrust-top-level-original:"; +const TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX = `\0${TOP_LEVEL_ORIGINAL_IMPORT_PREFIX}`; export interface LegacyBundlerPluginOptions { /** @@ -75,6 +84,18 @@ export const unplugin = createUnplugin( const allInstrumentations = getDefaultInstrumentationConfigs({ additionalInstrumentations: options.instrumentations, }); + const disabledIntegrationConfig = readDisabledInstrumentationEnvConfig( + process.env.BRAINTRUST_DISABLE_INSTRUMENTATION, + ).integrations; + const topLevelImportHookTarget: TopLevelImportHookTarget = + options.browser === true ? "browser" : "node"; + const topLevelImportHooks = getDefaultTopLevelImportHooks({ + disabledIntegrationConfig, + target: topLevelImportHookTarget, + }); + const originalSources = new Map(); + const originalDirectories = new Map(); + let nextOriginalId = 0; // Default to browser build, use polyfill unless explicitly disabled const dcModule = options.browser === false ? undefined : "dc-browser"; @@ -85,11 +106,37 @@ export const unplugin = createUnplugin( return { name: "code-transformer", enforce: "pre", + resolveId(id: string, importer?: string) { + if (id.startsWith(TOP_LEVEL_ORIGINAL_IMPORT_PREFIX)) { + return `${TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX}${id.slice( + TOP_LEVEL_ORIGINAL_IMPORT_PREFIX.length, + )}`; + } + if ( + id.startsWith(".") && + importer?.startsWith(TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX) + ) { + const originalDirectory = originalDirectories.get(importer); + if (originalDirectory) { + return resolve(originalDirectory, id); + } + } + return null; + }, + load(id: string) { + if (id.startsWith(TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX)) { + return originalSources.get(id) ?? null; + } + return null; + }, transform(code: string, id: string) { if (!id) { // Some modules apparently don't have an id? return null; } + if (id.startsWith(TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX)) { + return null; + } // Convert file:// URLs to regular paths at entry point // Node.js ESM loader hooks provide file:// URLs, but downstream code expects paths @@ -123,25 +170,52 @@ export const unplugin = createUnplugin( const normalizedModulePath = moduleDetails.path.replace(/\\/g, "/"); const moduleVersion = getModuleVersion(moduleDetails.basedir); + let nextCode = code; + let didPatch = false; + // Per-package source patches (see loader/special-case-patches.ts). - // Same anti-pattern fallback the runtime loader uses — mirrored here - // so bundled apps get the patches without relying on hook.mjs. - // Skipped for browser bundles since the wrapper templates use - // `node:module`/`require` to resolve `@mastra/observability`. + // Skipped for browser bundles to preserve the historical behavior of + // the OpenAI APIPromise patch path. if (options.browser !== true) { const patched = applySpecialCasePatch({ packageName: moduleName, modulePath: normalizedModulePath, - source: code, + source: nextCode, format: isModule ? "esm" : "cjs", }); if (patched !== null) { - return { code: patched, map: null }; + nextCode = patched; + didPatch = true; } } + const originalModuleSpecifier = `${TOP_LEVEL_ORIGINAL_IMPORT_PREFIX}${nextOriginalId++}`; + const originalModuleId = `${TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX}${originalModuleSpecifier.slice( + TOP_LEVEL_ORIGINAL_IMPORT_PREFIX.length, + )}`; + const topLevelWrapper = buildTopLevelImportHookSourceWrapper( + topLevelImportHooks, + { + format: isModule ? "esm" : "cjs", + modulePath: normalizedModulePath, + originalModuleSpecifier, + packageName: moduleName, + source: nextCode, + target: topLevelImportHookTarget, + }, + ); + if (topLevelWrapper !== null) { + originalSources.set(originalModuleId, nextCode); + originalDirectories.set(originalModuleId, dirname(filePath)); + nextCode = topLevelWrapper; + didPatch = true; + } + // If no version found if (!moduleVersion) { + if (didPatch) { + return { code: nextCode, map: null }; + } console.warn( `No 'package.json' version found for module ${moduleName} at ${moduleDetails.basedir}. Skipping transformation.`, ); @@ -157,13 +231,13 @@ export const unplugin = createUnplugin( if (!transformer) { // No instrumentations match this file - return null; + return didPatch ? { code: nextCode, map: null } : null; } try { // Transform the code const moduleType = isModule ? "esm" : "cjs"; - const result = transformer.transform(code, moduleType); + const result = transformer.transform(nextCode, moduleType); const transformedCode = result.code.replace( /const \{tracingChannel: ([A-Za-z_$][\w$]*)\} = ([A-Za-z_$][\w$]*);/g, "const $1 = $2.tracingChannel;", diff --git a/js/src/auto-instrumentations/bundler/webpack-loader.ts b/js/src/auto-instrumentations/bundler/webpack-loader.ts index c218c21a6..bb72bb0e3 100644 --- a/js/src/auto-instrumentations/bundler/webpack-loader.ts +++ b/js/src/auto-instrumentations/bundler/webpack-loader.ts @@ -27,6 +27,15 @@ import { readFileSync } from "fs"; import moduleDetailsFromPath from "module-details-from-path"; import { getDefaultInstrumentationConfigs } from "../configs/all"; import { type LegacyBundlerPluginOptions } from "./plugin"; +import { applySpecialCasePatch } from "../loader/special-case-patches"; +import { + buildTopLevelImportHookSourceWrapper, + getDefaultTopLevelImportHooks, + type TopLevelImportHookTarget, +} from "../loader/top-level-export-patches"; +import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config"; + +const TOP_LEVEL_ORIGINAL_QUERY = "braintrust-top-level-original"; /** * Helper function to get module version from package.json @@ -94,11 +103,15 @@ function codeTransformerLoader( const callback = this.async(); const options: LegacyBundlerPluginOptions = this.getOptions() ?? {}; const resourcePath: string = this.resourcePath; + const resourceQuery: string = this.resourceQuery ?? ""; // Skip virtual modules (e.g. Next.js loaders pass query-string URLs with no real path) if (!resourcePath) { return callback(null, code, inputSourceMap); } + if (resourceQuery.includes(TOP_LEVEL_ORIGINAL_QUERY)) { + return callback(null, code, inputSourceMap); + } // Determine if this is an ES module using multiple methods for accurate detection const ext = extname(resourcePath); @@ -122,12 +135,54 @@ function codeTransformerLoader( const moduleName = moduleDetails.name; const moduleVersion = getModuleVersion(moduleDetails.basedir); - if (!moduleVersion) { - return callback(null, code, inputSourceMap); - } - // Normalize the module path for Windows compatibility (WASM transformer expects forward slashes) const normalizedModulePath = moduleDetails.path.replace(/\\/g, "/"); + const moduleType: ModuleType = isModule ? "esm" : "cjs"; + const target: TopLevelImportHookTarget = + options.browser === true ? "browser" : "node"; + const disabledIntegrationConfig = readDisabledInstrumentationEnvConfig( + process.env.BRAINTRUST_DISABLE_INSTRUMENTATION, + ).integrations; + const topLevelImportHooks = getDefaultTopLevelImportHooks({ + disabledIntegrationConfig, + target, + }); + + let nextCode = code; + let didPatch = false; + + if (options.browser !== true) { + const patched = applySpecialCasePatch({ + format: moduleType, + modulePath: normalizedModulePath, + packageName: moduleName, + source: nextCode, + }); + if (patched !== null) { + nextCode = patched; + didPatch = true; + } + } + + const topLevelWrapper = buildTopLevelImportHookSourceWrapper( + topLevelImportHooks, + { + format: moduleType, + modulePath: normalizedModulePath, + originalModuleSpecifier: `${resourcePath}?${TOP_LEVEL_ORIGINAL_QUERY}`, + packageName: moduleName, + source: nextCode, + target, + }, + ); + if (topLevelWrapper !== null) { + nextCode = topLevelWrapper; + didPatch = true; + } + + if (!moduleVersion) { + return callback(null, nextCode, inputSourceMap); + } const matcher = getMatcher(options); const transformer = matcher.getTransformer( @@ -137,19 +192,18 @@ function codeTransformerLoader( ); if (!transformer) { - return callback(null, code, inputSourceMap); + return callback(null, nextCode, inputSourceMap); } try { - const moduleType: ModuleType = isModule ? "esm" : "cjs"; - const result = transformer.transform(code, moduleType); + const result = transformer.transform(nextCode, moduleType); callback(null, result.code, result.map ?? undefined); } catch (error) { console.warn( `[code-transformer-loader] Error transforming ${resourcePath}:`, error, ); - callback(null, code, inputSourceMap); + callback(null, didPatch ? nextCode : code, inputSourceMap); } } diff --git a/js/src/auto-instrumentations/configs/all.ts b/js/src/auto-instrumentations/configs/all.ts index b4642974a..f28905c91 100644 --- a/js/src/auto-instrumentations/configs/all.ts +++ b/js/src/auto-instrumentations/configs/all.ts @@ -96,12 +96,9 @@ const defaultInstrumentationConfigGroups: readonly InstrumentationConfigGroup[] configs: flueConfigs, }, // Note: `@mastra/core` is not listed here because its instrumentation - // doesn't go through the AST `code-transformer` matcher — Mastra's - // content-hashed chunks make `filePath`-based matching too brittle. - // Instead it's handled by the source-replacement entry in - // `loader/special-case-patches.ts`, which both the runtime loader - // (`hook.mjs` → `cjs-patch.ts`/`esm-hook.mts`) and the bundler plugin - // (`bundler/plugin.ts`) call. The `mastra` env-var disable still works. + // doesn't go through the AST `code-transformer` matcher. It is handled by + // the internal top-level import hook registry in + // `loader/top-level-export-patches.ts`. ]; export function getDefaultInstrumentationConfigs({ diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index 10626cebd..fa65e4a29 100644 --- a/js/src/auto-instrumentations/hook.mts +++ b/js/src/auto-instrumentations/hook.mts @@ -13,7 +13,7 @@ * - CJS modules: Transformed via ModulePatch monkey-patching Module._compile */ -import { register } from "node:module"; +import { register as registerModule } from "node:module"; import { isInstrumentationIntegrationDisabled, readDisabledInstrumentationEnvConfig, @@ -23,6 +23,28 @@ import { installMastraExporterFactory } from "./loader/mastra-observability-patc import { getDefaultAutoInstrumentationConfigs } from "./configs/all.js"; import { ModulePatch } from "./loader/cjs-patch.js"; import { patchTracingChannel } from "./patch-tracing-channel.js"; +import registerState from "./import-in-the-middle/lib/register.js"; +import { createHook as createImportInTheMiddleHook } from "./import-in-the-middle/create-hook.mjs"; +import { + getDefaultTopLevelImportHooks, + installTopLevelImportHookRunner, +} from "./loader/top-level-export-patches.js"; +import { installNodeTopLevelExportPatches } from "./loader/top-level-export-patches-node.js"; + +const BRAINTRUST_IITM_LOADER_PARAM = "braintrust-iitm-loader"; +const registryImportUrl = getCanonicalHookUrl(import.meta.url); +const asyncImportHookUrl = getImportInTheMiddleLoaderUrl(registryImportUrl); +const isImportInTheMiddleLoader = hasImportInTheMiddleLoaderParam( + import.meta.url, +); +const importInTheMiddleHook = createImportInTheMiddleHook(import.meta, { + registerUrl: registryImportUrl, +}); + +export const initialize = importInTheMiddleHook.initialize; +export const resolve = importInTheMiddleHook.resolve; +export const load = importInTheMiddleHook.load; +export const register = registerState.register; const state = ((globalThis as any)[ Symbol.for("braintrust.applyAutoInstrumentation") @@ -32,13 +54,13 @@ const alreadyApplied = state.applied; // Patch diagnostics_channel.tracePromise to handle APIPromise correctly. // MUST be done here (before any SDK code runs) to fix Anthropic APIPromise incompatibility. // Construct the module path dynamically to prevent build from stripping "node:" prefix. -if (!alreadyApplied) { +if (!isImportInTheMiddleLoader && !alreadyApplied) { const dcPath = ["node", "diagnostics_channel"].join(":"); const dc: any = await import(/* @vite-ignore */ dcPath as any); patchTracingChannel(dc.tracingChannel); } -if (!alreadyApplied) { +if (!isImportInTheMiddleLoader && !alreadyApplied) { const allConfigs = getDefaultAutoInstrumentationConfigs(); // Expose the Mastra exporter factory on globalThis so the loader patches @@ -52,8 +74,19 @@ if (!alreadyApplied) { installMastraExporterFactory(() => new BraintrustObservabilityExporter()); } + const topLevelImportHooks = getDefaultTopLevelImportHooks({ + disabledIntegrationConfig: disabled, + target: "node", + }); + installTopLevelImportHookRunner(topLevelImportHooks); + installNodeTopLevelExportPatches({ + asyncImportHookUrl, + hooks: topLevelImportHooks, + registryImportUrl, + }); + // 1. Register ESM loader for ESM modules - register("./loader/esm-hook.mjs", { + registerModule("./loader/esm-hook.mjs", { parentURL: import.meta.url, data: { instrumentations: allConfigs }, } as any); @@ -82,3 +115,24 @@ if (!alreadyApplied) { } } } + +function hasImportInTheMiddleLoaderParam(url: string): boolean { + try { + return new URL(url).searchParams.has(BRAINTRUST_IITM_LOADER_PARAM); + } catch { + return false; + } +} + +function getCanonicalHookUrl(url: string): string { + const parsed = new URL(url); + parsed.searchParams.delete(BRAINTRUST_IITM_LOADER_PARAM); + parsed.hash = ""; + return parsed.href; +} + +function getImportInTheMiddleLoaderUrl(canonicalUrl: string): string { + const parsed = new URL(canonicalUrl); + parsed.searchParams.set(BRAINTRUST_IITM_LOADER_PARAM, "true"); + return parsed.href; +} diff --git a/js/src/auto-instrumentations/import-in-the-middle/create-hook.mjs b/js/src/auto-instrumentations/import-in-the-middle/create-hook.mjs index 3ae4d53c3..778242514 100644 --- a/js/src/auto-instrumentations/import-in-the-middle/create-hook.mjs +++ b/js/src/auto-instrumentations/import-in-the-middle/create-hook.mjs @@ -292,9 +292,10 @@ function addIitm(url) { return urlObj.href; } -export function createHook(meta) { +export function createHook(meta, options = {}) { let cachedResolve; - const iitmURL = new URL("lib/register.js", meta.url).toString(); + const iitmURL = + options.registerUrl ?? new URL("lib/register.js", meta.url).toString(); const includeModules = new Set(); // Track CJS module URLs that IITM has wrapped. On Node 24+, CJS modules loaded diff --git a/js/src/auto-instrumentations/import-in-the-middle/create-hook.mjs.d.ts b/js/src/auto-instrumentations/import-in-the-middle/create-hook.mjs.d.ts new file mode 100644 index 000000000..a762ab4f9 --- /dev/null +++ b/js/src/auto-instrumentations/import-in-the-middle/create-hook.mjs.d.ts @@ -0,0 +1,23 @@ +export interface ImportInTheMiddleHook { + applyOptions(data: unknown): void; + initialize(data?: unknown): Promise; + load(url: string, context: unknown, parentLoad: Function): Promise; + loadSync(url: string, context: unknown, nextLoad: Function): unknown; + resolve( + specifier: string, + context: unknown, + parentResolve: Function, + ): Promise; + resolveSync( + specifier: string, + context: unknown, + nextResolve: Function, + ): unknown; +} + +export function createHook( + meta: { url: string }, + options?: { registerUrl?: string }, +): ImportInTheMiddleHook; + +export function supportsSyncHooks(): boolean; diff --git a/js/src/auto-instrumentations/import-in-the-middle/lib/get-exports.mjs b/js/src/auto-instrumentations/import-in-the-middle/lib/get-exports.mjs index 70920ce77..58d9b2b9b 100644 --- a/js/src/auto-instrumentations/import-in-the-middle/lib/get-exports.mjs +++ b/js/src/auto-instrumentations/import-in-the-middle/lib/get-exports.mjs @@ -154,6 +154,13 @@ const BUILT_INS = new Map(); let require; +function getRequire() { + if (!require) { + require = createRequire(import.meta.url); + } + return require; +} + // Returns a builtin's exports object. `process.getBuiltinModule` (Node >= // 20.16 / >= 22.3) bypasses registered loader hooks; `require` does not. Under // the in-thread `module.registerHooks` loader a plain `require(name)` here @@ -165,10 +172,7 @@ function loadBuiltin(name) { if (typeof process.getBuiltinModule === "function") { return process.getBuiltinModule(name); } - if (!require) { - require = createRequire(import.meta.url); - } - return require(name); + return getRequire()(name); } function getExportsForNodeBuiltIn(name) { @@ -284,11 +288,8 @@ function* getCjsExports(url, context, source) { [reUrl, reSpecifier] = resolved; } - if (!require) { - require = createRequire(import.meta.url); - } const newUrl = pathToFileURL( - require.resolve(reSpecifier, { + getRequire().resolve(reSpecifier, { paths: [dirname(fileURLToPath(reUrl))], }), ).href; diff --git a/js/src/auto-instrumentations/import-in-the-middle/lib/register.d.ts b/js/src/auto-instrumentations/import-in-the-middle/lib/register.d.ts new file mode 100644 index 000000000..229687dc8 --- /dev/null +++ b/js/src/auto-instrumentations/import-in-the-middle/lib/register.d.ts @@ -0,0 +1,23 @@ +export type Namespace = Record; + +export function register( + name: string, + namespace: Namespace, + set: Record boolean>, + get: Record unknown>, + specifier?: string, +): void; + +declare const registerState: { + addHookedModules(modules: string[]): void; + deleteHookedModules(modules: string[]): void; + hookedModules: Set; + importHooks: Array< + (name: string, namespace: Namespace, specifier?: string) => void + >; + register: typeof register; + specifiers: Map; + toHook: Array<[name: string, namespace: Namespace, specifier?: string]>; +}; + +export default registerState; diff --git a/js/src/auto-instrumentations/import-in-the-middle/lib/register.js b/js/src/auto-instrumentations/import-in-the-middle/lib/register.js index 08eecf643..34b2670a5 100644 --- a/js/src/auto-instrumentations/import-in-the-middle/lib/register.js +++ b/js/src/auto-instrumentations/import-in-the-middle/lib/register.js @@ -2,13 +2,27 @@ // // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2021 Datadog, Inc. -const importHooks = []; // TODO should this be a Set? -const setters = new WeakMap(); -const getters = new WeakMap(); -const specifiers = new Map(); -const toHook = []; -const hookedModuleCounts = new Map(); -const hookedModules = new Set(); +const state = (globalThis[ + Symbol.for("braintrust.importInTheMiddle.registerState") +] ??= { + getters: new WeakMap(), + hookedModuleCounts: new Map(), + hookedModules: new Set(), + importHooks: [], // TODO should this be a Set? + setters: new WeakMap(), + specifiers: new Map(), + toHook: [], +}); + +const { + getters, + hookedModuleCounts, + hookedModules, + importHooks, + setters, + specifiers, + toHook, +} = state; const proxyHandler = { set(target, name, value) { diff --git a/js/src/auto-instrumentations/loader/import-in-the-middle-runtime.ts b/js/src/auto-instrumentations/loader/import-in-the-middle-runtime.ts new file mode 100644 index 000000000..a19f82d87 --- /dev/null +++ b/js/src/auto-instrumentations/loader/import-in-the-middle-runtime.ts @@ -0,0 +1,158 @@ +import moduleDetailsFromPath from "module-details-from-path"; +import { isBuiltin } from "node:module"; +import { fileURLToPath } from "node:url"; +import registerState from "../import-in-the-middle/lib/register.js"; + +type Namespace = Record; +type HookFn = (exported: Namespace, name: string, baseDir?: string) => unknown; + +const { + addHookedModules, + deleteHookedModules, + importHooks, + specifiers, + toHook, +} = registerState; + +function addHook( + hook: (name: string, namespace: Namespace, specifier?: string) => void, +) { + importHooks.push(hook); + toHook.forEach(([name, namespace, specifier]) => + hook(name, namespace, specifier), + ); +} + +function removeHook( + hook: (name: string, namespace: Namespace, specifier?: string) => void, +) { + const index = importHooks.indexOf(hook); + if (index > -1) { + importHooks.splice(index, 1); + } +} + +function callHookFn( + hookFn: HookFn, + namespace: Namespace, + name: string, + baseDir?: string, +): void { + const newDefault = hookFn(namespace, name, baseDir); + if (newDefault && newDefault !== namespace && "default" in namespace) { + namespace.default = newDefault; + } +} + +function normalizeModules(modules: string[]): string[] { + if (!Array.isArray(modules) || modules.length === 0) { + throw new TypeError( + "Braintrust import-in-the-middle requires a non-empty modules array", + ); + } + + for (const each of modules) { + if (typeof each !== "string") { + throw new TypeError( + "Braintrust import-in-the-middle only supports string module names or file URLs", + ); + } + } + + return modules; +} + +function moduleMatches( + matchArg: string, + name: string, + filePath: string | undefined, + specifier: string | undefined, + baseDir: string | undefined, + loadUrl: string, +): boolean { + if (filePath && matchArg === filePath) { + return true; + } + + if (matchArg === specifier || matchArg === loadUrl || matchArg === name) { + return true; + } + + if (!baseDir) { + return false; + } + + // Keep the top-level package check from upstream, but do not support the + // broad internals mode. Internal files must be listed explicitly. + return matchArg === name && baseDir.endsWith(specifiers.get(loadUrl) ?? ""); +} + +export default class Hook { + private readonly modules: string[]; + private readonly iitmHook: ( + name: string, + namespace: Namespace, + specifier?: string, + ) => void; + + constructor(modules: string[], hookFn: HookFn) { + modules = normalizeModules(modules); + if (typeof hookFn !== "function") { + throw new TypeError( + "Braintrust import-in-the-middle requires a hook function", + ); + } + + addHookedModules(modules); + + this.modules = modules; + this.iitmHook = (name, namespace, specifier) => { + const loadUrl = name; + let filePath: string | undefined; + let baseDir: string | undefined; + + if (loadUrl.startsWith("node:")) { + const unprefixed = name.slice(5); + if (isBuiltin(unprefixed)) { + name = unprefixed; + } + } else if (loadUrl.startsWith("file://")) { + const stackTraceLimit = Error.stackTraceLimit; + Error.stackTraceLimit = 0; + try { + filePath = fileURLToPath(name); + name = filePath; + } catch {} + Error.stackTraceLimit = stackTraceLimit; + + if (filePath) { + const details = moduleDetailsFromPath(filePath); + if (details) { + name = details.name; + baseDir = details.basedir; + } + } + } + + for (const matchArg of modules) { + if ( + moduleMatches(matchArg, name, filePath, specifier, baseDir, loadUrl) + ) { + callHookFn( + hookFn, + namespace, + filePath && matchArg === filePath ? filePath : matchArg, + baseDir, + ); + } + } + }; + + addHook(this.iitmHook); + } + + unhook(): void { + removeHook(this.iitmHook); + deleteHookedModules(this.modules); + } +} diff --git a/js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts b/js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts index cda12afdf..0abc7d319 100644 --- a/js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts +++ b/js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts @@ -1,117 +1,176 @@ import { describe, expect, it } from "vitest"; import { - classifyMastraTarget, - patchMastraSource, + installMastraExporterFactory, + patchMastraExports, } from "./mastra-observability-patch"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; -describe("classifyMastraTarget", () => { - it("identifies @mastra/core main and submodule entries", () => { - expect(classifyMastraTarget("@mastra/core", "dist/index.js")).toBe("core"); - expect(classifyMastraTarget("@mastra/core", "dist/index.cjs")).toBe("core"); - expect(classifyMastraTarget("@mastra/core", "dist/mastra/index.js")).toBe( - "core", - ); - expect(classifyMastraTarget("@mastra/core", "dist/mastra/index.cjs")).toBe( - "core", - ); - }); +const fixturesDir = join( + dirname(fileURLToPath(import.meta.url)), + "../../../tests/auto-instrumentations/fixtures", +); +const mastraCoreBaseDir = join(fixturesDir, "node_modules/@mastra/core"); - it("identifies @mastra/observability entry", () => { - expect(classifyMastraTarget("@mastra/observability", "dist/index.js")).toBe( - "observability", - ); - expect( - classifyMastraTarget("@mastra/observability", "dist/index.cjs"), - ).toBe("observability"); - }); +describe("patchMastraExports — runtime @mastra/observability export", () => { + it("wraps Observability and appends the Braintrust exporter", () => { + delete (globalThis as any).__braintrustMastraExporterFactory; + installMastraExporterFactory(() => ({ name: "braintrust" })); - it("returns null for unrelated paths", () => { - expect( - classifyMastraTarget("@mastra/core", "dist/agent/index.js"), - ).toBeNull(); - expect( - classifyMastraTarget("@mastra/core", "dist/chunk-XYZ.js"), - ).toBeNull(); - expect(classifyMastraTarget("openai", "dist/index.js")).toBeNull(); - }); -}); + class Observability { + public config: unknown; + constructor(config: unknown) { + this.config = config; + } + } + const namespace = { Observability }; + + patchMastraExports(namespace, { moduleName: "@mastra/observability" }); + patchMastraExports(namespace, { moduleName: "@mastra/observability" }); -describe("patchMastraSource — @mastra/core ESM entry", () => { - it("rewrites a thin re-export into a Proxy-wrapped class", () => { - const original = `export { Mastra } from './chunk-PLCLLPJL.js';\n`; - const patched = patchMastraSource(original, "core", "esm"); - - // The rewritten source must import the original class from the same chunk - expect(patched).toContain( - `import { Mastra as __braintrustOrigMastra } from "./chunk-PLCLLPJL.js"`, - ); - // It must wrap in a Proxy with a `construct` trap - expect(patched).toContain("new Proxy(__braintrustOrigMastra,"); - expect(patched).toContain("construct(target, args, newTarget)"); - // It must re-export `Mastra` so consumers' bindings are unchanged - expect(patched).toContain("export { Mastra }"); - // It must pull Observability via createRequire so it resolves from the - // user's node_modules tree (not our SDK's) - expect(patched).toContain("createRequire"); - expect(patched).toContain(`__braintrustRequire("@mastra/observability")`); + const instance = new namespace.Observability({ + configs: { + default: { + exporters: [{ name: "other" }], + serviceName: "service", + }, + }, + }); + + expect(instance.config).toEqual({ + configs: { + default: { + exporters: [{ name: "other" }, { name: "braintrust" }], + serviceName: "service", + }, + }, + }); }); - it("returns the original source when the re-export shape doesn't match", () => { - const original = `// arbitrary code that doesn't re-export Mastra\n`; - const patched = patchMastraSource(original, "core", "esm"); - expect(patched).toBe(original); + it("does not duplicate an existing Braintrust exporter", () => { + delete (globalThis as any).__braintrustMastraExporterFactory; + installMastraExporterFactory(() => ({ name: "braintrust" })); + + class Observability { + public config: unknown; + constructor(config: unknown) { + this.config = config; + } + } + const namespace = { Observability }; + + patchMastraExports(namespace, { moduleName: "@mastra/observability" }); + + const instance = new namespace.Observability({ + configs: { + default: { + exporters: [{ name: "braintrust" }], + serviceName: "service", + }, + }, + }); + + expect(instance.config).toEqual({ + configs: { + default: { + exporters: [{ name: "braintrust" }], + serviceName: "service", + }, + }, + }); }); - it("preserves the exact chunk path Mastra references", () => { - const original = `export { Mastra } from '../chunk-DIFFERENT.js';\n`; - const patched = patchMastraSource(original, "core", "esm"); - expect(patched).toContain("../chunk-DIFFERENT.js"); + it("creates a default config when none is provided", () => { + delete (globalThis as any).__braintrustMastraExporterFactory; + installMastraExporterFactory(() => ({ name: "braintrust" })); + + class Observability { + public config: unknown; + constructor(config: unknown) { + this.config = config; + } + } + const namespace = { Observability }; + + patchMastraExports(namespace, { moduleName: "@mastra/observability" }); + + const instance = new namespace.Observability(undefined); + + expect(instance.config).toEqual({ + configs: { + default: { + exporters: [{ name: "braintrust" }], + serviceName: "mastra", + }, + }, + }); }); }); -describe("patchMastraSource — @mastra/core CJS entry", () => { - it("rewrites the require + defineProperty shape into a Proxy", () => { - const original = `'use strict'; -var chunkVOP4TUHG_cjs = require('./chunk-VOP4TUHG.cjs'); -Object.defineProperty(exports, "Mastra", { - enumerable: true, - get: function () { return chunkVOP4TUHG_cjs.Mastra; } -}); -`; - const patched = patchMastraSource(original, "core", "cjs"); - - expect(patched).toContain(`require("./chunk-VOP4TUHG.cjs")`); - expect(patched).toContain(`require("@mastra/observability")`); - expect(patched).toContain("__braintrustChunk.Mastra"); - expect(patched).toContain("new Proxy("); - expect(patched).toContain(`Object.defineProperty(exports, "Mastra"`); +describe("patchMastraExports — runtime @mastra/core export", () => { + it("wraps Mastra and injects default Observability when config has none", () => { + delete (globalThis as any).__braintrustMastraExporterFactory; + installMastraExporterFactory(() => ({ name: "braintrust" })); + + class Mastra { + public observability: any; + constructor(config: any = {}) { + this.observability = config.observability; + } + } + const namespace = { Mastra }; + + patchMastraExports(namespace, { + baseDir: mastraCoreBaseDir, + moduleName: "@mastra/core", + }); + patchMastraExports(namespace, { + baseDir: mastraCoreBaseDir, + moduleName: "@mastra/core", + }); + + const instance = new namespace.Mastra({}); + + expect(instance.observability).toBeTruthy(); + expect( + instance.observability.config.configs.default.exporters.map( + (exporter: { name: string }) => exporter.name, + ), + ).toEqual(["braintrust"]); }); -}); -describe("patchMastraSource — @mastra/observability entry", () => { - it("appends a Proxy wrap, leaving the original source intact above", () => { - const original = `var Observability = class extends MastraBase {};\nexport { Observability };\n`; - const patched = patchMastraSource(original, "observability", "esm"); - - // Original source still starts the file - expect(patched.startsWith(original)).toBe(true); - // Append wraps the Observability binding via Proxy - expect(patched).toContain("function __braintrustWrapObservability"); - expect(patched).toContain( - `if (typeof Observability === "undefined") return`, - ); - expect(patched).toContain("new Proxy(__OriginalObservability"); - expect(patched).toContain("factory()"); - expect(patched).toContain("__braintrustWrapped"); + it("preserves a user-provided observability config", () => { + delete (globalThis as any).__braintrustMastraExporterFactory; + installMastraExporterFactory(() => ({ name: "braintrust" })); + + class Mastra { + public observability: any; + constructor(config: any = {}) { + this.observability = config.observability; + } + } + const namespace = { Mastra }; + const observability = { user: true }; + + patchMastraExports(namespace, { + baseDir: mastraCoreBaseDir, + moduleName: "@mastra/core", + }); + + const instance = new namespace.Mastra({ observability }); + + expect(instance.observability).toBe(observability); }); - it("doesn't depend on chunk path extraction for observability", () => { - // The Observability entry is one big inline file, not a re-export. - // patchMastraSource should still produce a valid append even when no - // chunk-shaped pattern is present in the source. - const original = `// big inline bundle\nvar Observability = class {};\nexport { Observability };\n`; - const patched = patchMastraSource(original, "observability", "esm"); - expect(patched).not.toBe(original); - expect(patched.length).toBeGreaterThan(original.length); + it("is a no-op when @mastra/observability is unavailable", () => { + class Mastra {} + const namespace = { Mastra }; + + patchMastraExports(namespace, { + baseDir: dirname(fileURLToPath(import.meta.url)), + moduleName: "@mastra/core", + }); + + expect(namespace.Mastra).toBe(Mastra); }); }); diff --git a/js/src/auto-instrumentations/loader/mastra-observability-patch.ts b/js/src/auto-instrumentations/loader/mastra-observability-patch.ts index 578b7a48c..22879923e 100644 --- a/js/src/auto-instrumentations/loader/mastra-observability-patch.ts +++ b/js/src/auto-instrumentations/loader/mastra-observability-patch.ts @@ -1,46 +1,14 @@ /** - * Runtime patches that auto-install the Braintrust observability exporter for - * Mastra users. + * Runtime hook implementation for Mastra top-level exports. * - * Two entry shapes need handling, and they're each different enough that we - * pattern-match them separately rather than trying to share one strategy: - * - * 1. **`@mastra/core` `Mastra` entries** (`dist/index.{js,cjs}` and - * `dist/mastra/index.{js,cjs}`) are thin re-exports from a content-hashed - * chunk that changes filename every release: - * - * // dist/index.js - * export { Mastra } from "./chunk-PLCLLPJL.js"; - * - * The entry filename itself is stable (pinned by the package's `exports` - * map), so we patch it — but a re-export has no local binding to mutate, - * so we have to rewrite the line into an explicit import + Proxy + export. - * That means extracting the chunk path with a regex. The regex pattern is - * stable; only the chunk hash inside the matched string changes. - * - * 2. **`@mastra/observability` `Observability` entry** (`dist/index.{js,cjs}`) - * inlines its class declaration in the entry itself: - * - * var Observability = class extends MastraBase { ... }; - * export { ..., Observability, ... }; - * - * Because the local `Observability` is a `var` binding, we can *append* a - * wrap (reassign `Observability` to a Proxy) and ESM's live-binding - * semantics propagate it to importers. No source rewrite needed, no chunk - * path involved. - * - * In both cases the generated wrapper is a `Proxy` with a `construct` trap. - * On the Mastra side the trap injects an `Observability` instance into the - * config when the user didn't pass one (via a factory registered on - * `globalThis` by `hook.mjs`). On the Observability side the trap walks the - * `configs` map and appends our exporter to any instance config that doesn't - * already have one with `name === "braintrust"`. - * - * The patch is a no-op if the regex doesn't match (e.g., Mastra restructures - * its build) or if our globalThis factory isn't present (e.g., user disabled - * via `BRAINTRUST_DISABLE_INSTRUMENTATION=mastra`). + * The IITM/RITM adapters pass real module export objects here. Bundler/source + * adapters generate wrapper modules that build the same mutable namespace + * facade, call the shared hook runner, and re-export the updated bindings. */ +import { createRequire } from "node:module"; +import { join } from "node:path"; + /** * Name of the `globalThis` property the generated patches read to look up the * Braintrust exporter factory at runtime. Set by `hook.mjs` (loader path) and @@ -65,224 +33,247 @@ export function installMastraExporterFactory(factory: () => unknown): void { const MASTRA_CORE_PACKAGE = "@mastra/core"; const MASTRA_OBSERVABILITY_PACKAGE = "@mastra/observability"; +const MASTRA_CORE_MASTRA_SPECIFIER = "@mastra/core/mastra"; +const MASTRA_RUNTIME_WRAPPED = Symbol.for( + "braintrust.mastra.runtime-export-wrapped", +); -// Entrypoints pinned by each package's `exports` map. Both Mastra entries -// (`dist/index.js` and `dist/mastra/index.js`) re-export `Mastra` from a -// chunk, so they both need the rewrite treatment. -const MASTRA_CORE_ENTRY_PATHS = new Set([ - "dist/index.js", - "dist/index.cjs", - "dist/mastra/index.js", - "dist/mastra/index.cjs", -]); -const MASTRA_OBSERVABILITY_ENTRY_PATHS = new Set([ - "dist/index.js", - "dist/index.cjs", -]); +export interface MastraRuntimePatchContext { + moduleName: string; + baseDir?: string; + resolutionBase?: string; +} -export type MastraTargetKind = "core" | "observability"; -export type MastraModuleFormat = "esm" | "cjs"; +type RuntimeConstructor = new (...args: unknown[]) => unknown; -export function classifyMastraTarget( - packageName: string, - modulePath: string, -): MastraTargetKind | null { +export function patchMastraExports( + exportsValue: T, + context: MastraRuntimePatchContext, +): T { if ( - packageName === MASTRA_CORE_PACKAGE && - MASTRA_CORE_ENTRY_PATHS.has(modulePath) + context.moduleName === MASTRA_CORE_PACKAGE || + context.moduleName === MASTRA_CORE_MASTRA_SPECIFIER ) { - return "core"; + return patchMastraCoreExports(exportsValue, context) as T; + } else if (context.moduleName === MASTRA_OBSERVABILITY_PACKAGE) { + return patchMastraObservabilityExports(exportsValue) as T; } - if ( - packageName === MASTRA_OBSERVABILITY_PACKAGE && - MASTRA_OBSERVABILITY_ENTRY_PATHS.has(modulePath) - ) { - return "observability"; + return exportsValue; +} + +function patchMastraCoreExports( + exportsValue: unknown, + context: MastraRuntimePatchContext, +): unknown { + if (!exportsValue || typeof exportsValue !== "object") return exportsValue; + + const namespace = exportsValue as Record; + const Mastra = asRuntimeConstructor(namespace.Mastra); + if (!Mastra || isRuntimeWrapped(Mastra)) return exportsValue; + + const Observability = loadObservabilityClass(context); + if (!Observability) return exportsValue; + + const wrapped = new Proxy(Mastra, { + construct(target, args, newTarget) { + const firstArg = args[0]; + if ( + (!firstArg || + typeof firstArg !== "object" || + !("observability" in firstArg) || + !(firstArg as { observability?: unknown }).observability) && + Observability + ) { + try { + const observability = new Observability({ + configs: { default: { serviceName: "mastra" } }, + }); + const nextConfig = + firstArg && typeof firstArg === "object" + ? { ...(firstArg as Record), observability } + : { observability }; + return Reflect.construct( + target, + [nextConfig, ...args.slice(1)], + newTarget, + ); + } catch { + // Mastra will keep its own fallback behavior if constructing + // Observability fails. + } + } + return Reflect.construct(target, args, newTarget); + }, + }); + markRuntimeWrapped(wrapped); + + if (setNamespaceExport(namespace, "Mastra", wrapped)) { + return exportsValue; } - return null; + return cloneNamespaceWithExport(namespace, "Mastra", wrapped); } -// Extract the chunk path from a Mastra entry source. -// ESM shape: `export { Mastra } from "../chunk-XYZ.js";` -// CJS shape: `var chunkXYZ_cjs = require("../chunk-XYZ.cjs"); ...` -// Returns null when the shape isn't what we expect — caller falls through and -// emits the original source unchanged so user code keeps working. -function extractChunkPath(source: string): string | null { - const esmMatch = source.match( - /export\s*\{\s*Mastra(?:\s+as\s+\w+)?\s*\}\s*from\s*['"]([^'"]+)['"]/, - ); - if (esmMatch) return esmMatch[1]; +function patchMastraObservabilityExports(exportsValue: unknown): unknown { + if (!exportsValue || typeof exportsValue !== "object") return exportsValue; - const cjsMatch = source.match( - /require\s*\(\s*['"]([^'"]+chunk-[^'"]+)['"]\s*\)/, - ); - if (cjsMatch) return cjsMatch[1]; + const namespace = exportsValue as Record; + const Observability = asRuntimeConstructor(namespace.Observability); + if (!Observability || isRuntimeWrapped(Observability)) { + return exportsValue; + } + + const wrapped = new Proxy(Observability, { + construct(target, args, newTarget) { + const nextArgs = args.slice(); + nextArgs[0] = ensureBraintrustExporter(nextArgs[0]); + return Reflect.construct(target, nextArgs, newTarget); + }, + }); + markRuntimeWrapped(wrapped); - return null; + if (setNamespaceExport(namespace, "Observability", wrapped)) { + return exportsValue; + } + return cloneNamespaceWithExport(namespace, "Observability", wrapped); +} + +function loadObservabilityClass( + context: MastraRuntimePatchContext, +): RuntimeConstructor | undefined { + try { + const requireFromMastra = createRequire( + context.resolutionBase ?? + (context.baseDir + ? join(context.baseDir, "package.json") + : join(process.cwd(), "package.json")), + ); + const observability = requireFromMastra(MASTRA_OBSERVABILITY_PACKAGE); + patchMastraObservabilityExports(observability); + return asRuntimeConstructor(observability?.Observability); + } catch { + return undefined; + } } -// All wrapper templates below are emitted as runtime JS in the target module's -// scope. They never reference Braintrust SDK code directly — instead they look -// up a factory function on `globalThis` (registered by `hook.mjs`). That keeps -// the patch independent of how the user's module graph resolves `braintrust`, -// and lets us cleanly no-op when the user disables our integration via env. +function ensureBraintrustExporter(rawConfig: unknown): unknown { + try { + const factory = (globalThis as Record)[ + MASTRA_EXPORTER_FACTORY_GLOBAL + ]; + if (typeof factory !== "function") return rawConfig; -const EXPORTER_FACTORY_KEY = JSON.stringify(MASTRA_EXPORTER_FACTORY_GLOBAL); + const config = + rawConfig && typeof rawConfig === "object" + ? (rawConfig as Record) + : {}; + const configsIn = + config.configs && typeof config.configs === "object" + ? (config.configs as Record) + : undefined; + const configsOut: Record = {}; + let hadEntries = false; -// Construct-trap body shared between ESM and CJS wrappers. The wrapper loads -// `@mastra/observability` at module-evaluation time from the Mastra entry's -// own location (so the resolver finds the user's `node_modules`, not our -// SDK's). When the user constructs `new Mastra(...)` without an observability -// config, the trap injects a default `Observability` — which, because it -// loaded through our patched ESM hook / CJS patch, already auto-installs the -// Braintrust exporter via its own constructor wrap. -const MASTRA_PROXY_HANDLER_BODY = ` -{ - construct(target, args, newTarget) { - var firstArg = args[0]; - if ( - (!firstArg || typeof firstArg !== "object" || !firstArg.observability) && - __braintrustObservabilityClass - ) { - try { - // serviceName is required by Mastra's Observability validator; pass - // something sensible by default. Users who want a different name - // should construct Observability themselves. - var observability = new __braintrustObservabilityClass({ - configs: { default: { serviceName: "mastra" } }, - }); - args = args.slice(); - args[0] = Object.assign({}, firstArg, { observability: observability }); - } catch (e) { - // Fall through. Mastra will use its own NoOp; user code still works, - // just without auto-instrumentation. + if (configsIn) { + for (const [name, rawInstanceConfig] of Object.entries(configsIn)) { + hadEntries = true; + const instanceConfig = + rawInstanceConfig && typeof rawInstanceConfig === "object" + ? (rawInstanceConfig as Record) + : {}; + const existing = Array.isArray(instanceConfig.exporters) + ? instanceConfig.exporters + : []; + const hasOurs = existing.some( + (exporter) => + exporter && + typeof exporter === "object" && + (exporter as { name?: unknown }).name === "braintrust", + ); + configsOut[name] = { + ...instanceConfig, + exporters: hasOurs ? existing : [...existing, factory()], + }; } } - return Reflect.construct(target, args, newTarget); - }, -}`; -function buildMastraEsmWrapper(chunkPath: string): string { - const chunk = JSON.stringify(chunkPath); - return `import { Mastra as __braintrustOrigMastra } from ${chunk}; -import { createRequire as __braintrustCreateRequire } from "node:module"; + if (!hadEntries) { + configsOut.default = { + serviceName: "mastra", + exporters: [factory()], + }; + } -let __braintrustObservabilityClass = null; -try { - // Resolve @mastra/observability relative to this module (the Mastra entry), - // so it's looked up from the user's node_modules tree. - const __braintrustRequire = __braintrustCreateRequire(import.meta.url); - __braintrustObservabilityClass = - __braintrustRequire("@mastra/observability").Observability; -} catch (e) { - // @mastra/observability isn't installed; the construct trap will skip the - // auto-construct branch and Mastra falls back to its own NoOp. -} -const Mastra = new Proxy(__braintrustOrigMastra, ${MASTRA_PROXY_HANDLER_BODY}); -export { Mastra }; -`; + return { ...config, configs: configsOut }; + } catch { + return rawConfig; + } } -function buildMastraCjsWrapper(chunkPath: string): string { - const chunk = JSON.stringify(chunkPath); - return `'use strict'; -const __braintrustChunk = require(${chunk}); - -let __braintrustObservabilityClass = null; -try { - __braintrustObservabilityClass = require("@mastra/observability").Observability; -} catch (e) { - // @mastra/observability isn't installed; same fallback as the ESM wrapper. +function asRuntimeConstructor(value: unknown): RuntimeConstructor | undefined { + return typeof value === "function" + ? (value as RuntimeConstructor) + : undefined; } -const __braintrustWrappedMastra = new Proxy( - __braintrustChunk.Mastra, - ${MASTRA_PROXY_HANDLER_BODY}, -); -Object.defineProperty(exports, "Mastra", { - enumerable: true, - configurable: true, - get: function () { return __braintrustWrappedMastra; } -}); -`; +function isRuntimeWrapped(value: RuntimeConstructor): boolean { + return ( + (value as unknown as Record)[MASTRA_RUNTIME_WRAPPED] === + true + ); } -// Appended to the end of `@mastra/observability/dist/index.{js,cjs}`. The -// entry declares `var Observability = class extends MastraBase { ... }`, -// which means we can reassign the binding from inside the same module after -// the class is created and the export line has run. ESM live-binding -// semantics make external importers see the new value; CJS lookups go -// through `exports.Observability`, which we redefine to mirror the wrap. -const OBSERVABILITY_APPEND_BODY = ` -;(function __braintrustWrapObservability() { - // Top-level so we can both read and reassign the var binding the original - // entry declared. - if (typeof Observability === "undefined") return; - if (Observability.__braintrustWrapped) return; - function __braintrustEnsureExporter(rawConfig) { - try { - var factory = globalThis[${EXPORTER_FACTORY_KEY}]; - if (typeof factory !== "function") return rawConfig; - var config = rawConfig && typeof rawConfig === "object" ? rawConfig : {}; - var configsIn = config.configs && typeof config.configs === "object" ? config.configs : null; - var configsOut = {}; - var hadEntries = false; - if (configsIn) { - for (var name in configsIn) { - if (!Object.prototype.hasOwnProperty.call(configsIn, name)) continue; - hadEntries = true; - var inst = configsIn[name] || {}; - var existing = Array.isArray(inst.exporters) ? inst.exporters : []; - var hasOurs = existing.some(function (e) { return e && e.name === "braintrust"; }); - configsOut[name] = Object.assign({}, inst, { - exporters: hasOurs ? existing : existing.concat([factory()]), - }); - } - } - if (!hadEntries) { - configsOut.default = { - serviceName: "mastra", - exporters: [factory()], - }; - } - return Object.assign({}, config, { configs: configsOut }); - } catch (e) { - return rawConfig; - } +function markRuntimeWrapped(value: RuntimeConstructor): void { + try { + Object.defineProperty(value, MASTRA_RUNTIME_WRAPPED, { + configurable: false, + enumerable: false, + value: true, + }); + } catch { + // Best effort only. Idempotence still holds through the export assignment. } - var __OriginalObservability = Observability; - Observability = new Proxy(__OriginalObservability, { - construct: function (target, args, newTarget) { - var nextArgs = args.slice(); - nextArgs[0] = __braintrustEnsureExporter(nextArgs[0]); - return Reflect.construct(target, nextArgs, newTarget); - }, - }); - Observability.__braintrustWrapped = true; - if (typeof exports !== "undefined" && exports && typeof exports === "object") { - try { - Object.defineProperty(exports, "Observability", { - enumerable: true, - configurable: true, - get: function () { return Observability; }, - }); - } catch (e) {} +} + +function setNamespaceExport( + namespace: Record, + key: string, + value: RuntimeConstructor, +): boolean { + try { + namespace[key] = value; + if (namespace[key] === value) return true; + } catch { + // Try defineProperty below. } -})(); -`; -export function patchMastraSource( - source: string, - target: MastraTargetKind, - format: MastraModuleFormat, -): string { - if (target === "core") { - const chunkPath = extractChunkPath(source); - if (!chunkPath) return source; - return format === "esm" - ? buildMastraEsmWrapper(chunkPath) - : buildMastraCjsWrapper(chunkPath); + try { + Object.defineProperty(namespace, key, { + configurable: true, + enumerable: true, + value, + writable: true, + }); + return namespace[key] === value; + } catch { + return false; } - // Observability: append-wrap, leaving original source intact above. - return source + OBSERVABILITY_APPEND_BODY; +} + +function cloneNamespaceWithExport( + namespace: Record, + key: string, + value: RuntimeConstructor, +): Record { + return Object.defineProperties( + Object.create(Object.getPrototypeOf(namespace)), + { + ...Object.getOwnPropertyDescriptors(namespace), + [key]: { + configurable: true, + enumerable: true, + value, + writable: true, + }, + }, + ); } diff --git a/js/src/auto-instrumentations/loader/require-in-the-middle-runtime.ts b/js/src/auto-instrumentations/loader/require-in-the-middle-runtime.ts new file mode 100644 index 000000000..7dbc4d854 --- /dev/null +++ b/js/src/auto-instrumentations/loader/require-in-the-middle-runtime.ts @@ -0,0 +1,266 @@ +import moduleDetailsFromPath from "module-details-from-path"; +import Module, { builtinModules, createRequire } from "node:module"; +import { parse as parsePath } from "node:path"; +import { pathToFileURL } from "node:url"; + +type OnRequireFn = (exports: T, name: string, basedir?: string) => T; +type GetBuiltinModuleFn = (this: unknown, id: string) => unknown; +type ModuleWithInternals = typeof Module & { + _cache: Record) | undefined>; + _resolveFilename(id: string, parent: NodeJS.Module): string; +}; + +let builtinModulesSet: Set | undefined; +const ModuleInternals = Module as ModuleWithInternals; +const requireForResolve = createRequire( + pathToFileURL(process.argv[1] ?? process.cwd()).href, +); + +const isCore = + typeof Module.isBuiltin === "function" + ? Module.isBuiltin + : (moduleName: string) => { + if (moduleName.startsWith("node:")) { + return true; + } + + builtinModulesSet ??= new Set(builtinModules); + return builtinModulesSet.has(moduleName); + }; + +const normalize = /([/\\]index)?(\.js|\.cjs)?$/; + +class ExportsCache { + private readonly localCache = new Map(); + private readonly kRitmExports = Symbol("RitmExports"); + + has(filename: string, isBuiltin: boolean): boolean { + if (this.localCache.has(filename)) { + return true; + } else if (!isBuiltin) { + const mod = ModuleInternals._cache[filename] as + | (NodeJS.Module & Record) + | undefined; + return !!(mod && this.kRitmExports in mod); + } else { + return false; + } + } + + get(filename: string, isBuiltin: boolean): unknown { + const cachedExports = this.localCache.get(filename); + if (cachedExports !== undefined) { + return cachedExports; + } else if (!isBuiltin) { + const mod = ModuleInternals._cache[filename] as + | (NodeJS.Module & Record) + | undefined; + return mod && mod[this.kRitmExports]; + } + } + + set(filename: string, exports: unknown, isBuiltin: boolean): void { + if (isBuiltin) { + this.localCache.set(filename, exports); + } else if (filename in ModuleInternals._cache) { + ( + ModuleInternals._cache[filename] as NodeJS.Module & + Record + )[this.kRitmExports] = exports; + } else { + this.localCache.set(filename, exports); + } + } +} + +function normalizeModules(modules: string[]): string[] { + if (!Array.isArray(modules) || modules.length === 0) { + throw new TypeError( + "Braintrust require-in-the-middle requires a non-empty modules array", + ); + } + + for (const each of modules) { + if (typeof each !== "string") { + throw new TypeError( + "Braintrust require-in-the-middle only supports string module names or absolute paths", + ); + } + } + + return modules; +} + +export default class Hook { + private readonly cache = new ExportsCache(); + private readonly modules: string[]; + private readonly onrequire: OnRequireFn; + private unhooked = false; + private readonly origRequire = Module.prototype.require; + private readonly origGetBuiltinModule: GetBuiltinModuleFn | undefined = ( + process as unknown as { getBuiltinModule?: GetBuiltinModuleFn } + ).getBuiltinModule; + + constructor(modules: string[], onrequire: OnRequireFn) { + modules = normalizeModules(modules); + if (typeof onrequire !== "function") { + throw new TypeError( + "Braintrust require-in-the-middle requires an onrequire function", + ); + } + + if (typeof ModuleInternals._resolveFilename !== "function") { + throw new Error( + `Expected Module._resolveFilename to be a function, got ${typeof ModuleInternals._resolveFilename}`, + ); + } + + this.modules = modules; + this.onrequire = onrequire; + const self = this; + const patching = new Set(); + + Module.prototype.require = function (id: string) { + if (self.unhooked === true) { + return self.origRequire.apply(this, arguments as any); + } + + return patchedRequire.call(this, arguments, false); + }; + + if (typeof this.origGetBuiltinModule === "function") { + ( + process as unknown as { getBuiltinModule: GetBuiltinModuleFn } + ).getBuiltinModule = function (id: string) { + if (self.unhooked === true) { + return self.origGetBuiltinModule!.apply(process, arguments as any); + } + + return patchedRequire.call( + process as unknown as NodeJS.Module, + arguments, + true, + ); + }; + } + + function patchedRequire( + this: NodeJS.Module, + args: IArguments, + coreOnly: boolean, + ): unknown { + const id = args[0] as string; + const core = isCore(id); + let filename: string; + if (core) { + filename = id; + if (id.startsWith("node:")) { + const idWithoutPrefix = id.slice(5); + if (isCore(idWithoutPrefix)) { + filename = idWithoutPrefix; + } + } + } else if (coreOnly) { + return self.origGetBuiltinModule!.call(process, id); + } else { + try { + filename = ModuleInternals._resolveFilename(id, this); + } catch { + return self.origRequire.apply(this, args as any); + } + } + + if (self.cache.has(filename, core) === true) { + return self.cache.get(filename, core); + } + + const isPatching = patching.has(filename); + if (isPatching === false) { + patching.add(filename); + } + + const exports = coreOnly + ? self.origGetBuiltinModule!.call(process, id) + : self.origRequire.apply(this, args as any); + + if (isPatching === true) { + return exports; + } + + patching.delete(filename); + + let moduleName: string; + let basedir: string | undefined; + + if (core === true) { + if (modules.includes(filename) === false) { + return exports; + } + moduleName = filename; + } else if (modules.includes(filename)) { + const parsedPath = parsePath(filename); + moduleName = parsedPath.name; + basedir = parsedPath.dir; + } else { + const stat = moduleDetailsFromPath(filename); + if (!stat) { + return exports; + } + moduleName = stat.name; + basedir = stat.basedir; + + const fullModuleName = resolveModuleName(stat); + let matchFound = false; + if (!id.startsWith(".") && modules.includes(id)) { + moduleName = id; + matchFound = true; + } + + if ( + !modules.includes(moduleName) && + !modules.includes(fullModuleName) + ) { + return exports; + } + + if (modules.includes(fullModuleName) && fullModuleName !== moduleName) { + moduleName = fullModuleName; + matchFound = true; + } + + if (!matchFound) { + let res: string; + try { + res = requireForResolve.resolve(moduleName, { paths: [basedir] }); + } catch { + self.cache.set(filename, exports, core); + return exports; + } + + if (res !== filename) { + self.cache.set(filename, exports, core); + return exports; + } + } + } + + const patchedExports = onrequire(exports, moduleName, basedir); + self.cache.set(filename, patchedExports, core); + return patchedExports; + } + } + + unhook(): void { + this.unhooked = true; + Module.prototype.require = this.origRequire; + if (typeof this.origGetBuiltinModule === "function") { + ( + process as unknown as { getBuiltinModule: GetBuiltinModuleFn } + ).getBuiltinModule = this.origGetBuiltinModule; + } + } +} + +function resolveModuleName(stat: { name: string; path: string }): string { + return `${stat.name}/${stat.path.replace(normalize, "")}`; +} diff --git a/js/src/auto-instrumentations/loader/special-case-patches.ts b/js/src/auto-instrumentations/loader/special-case-patches.ts index 502b0af51..8f6e37728 100644 --- a/js/src/auto-instrumentations/loader/special-case-patches.ts +++ b/js/src/auto-instrumentations/loader/special-case-patches.ts @@ -21,23 +21,11 @@ * .parse()` doesn't double-read the response body. Removable once OpenAI * stops sharing the same `APIPromise` between `create()` and * `_thenUnwrap()`. - * - `@mastra/core` and `@mastra/observability` entries: Mastra ships - * code-split bundles with content-hashed chunk filenames, so we patch - * the stable submodule entries to install the - * `BraintrustObservabilityExporter` automatically. Removable when Mastra - * adopts a NPM-installable Braintrust exporter package directly, or when - * `import-in-the-middle` is reliable enough across Node versions to use - * for the same job. */ -import { - classifyMastraTarget, - patchMastraSource, - type MastraModuleFormat, -} from "./mastra-observability-patch.js"; import { OPENAI_API_PROMISE_PATCH } from "./openai-api-promise-patch.js"; -export type SpecialCaseFormat = MastraModuleFormat; +export type SpecialCaseFormat = "esm" | "cjs"; export interface SpecialCaseInput { packageName: string; @@ -62,16 +50,6 @@ export function applySpecialCasePatch(input: SpecialCaseInput): string | null { return input.source + OPENAI_API_PROMISE_PATCH; } - // Mastra: rewrite the stable submodule entries (@mastra/core) or append a - // Proxy wrap to the inline class binding (@mastra/observability). - const mastraTarget = classifyMastraTarget( - input.packageName, - input.modulePath, - ); - if (mastraTarget) { - return patchMastraSource(input.source, mastraTarget, input.format); - } - return null; } @@ -87,8 +65,5 @@ export function isSpecialCaseTarget( if (packageName === "openai" && modulePath.includes("api-promise")) { return true; } - if (classifyMastraTarget(packageName, modulePath) !== null) { - return true; - } return false; } diff --git a/js/src/auto-instrumentations/loader/top-level-export-patches-node.test.ts b/js/src/auto-instrumentations/loader/top-level-export-patches-node.test.ts new file mode 100644 index 000000000..6bb7ad72b --- /dev/null +++ b/js/src/auto-instrumentations/loader/top-level-export-patches-node.test.ts @@ -0,0 +1,184 @@ +import { describe, expect, it } from "vitest"; +import { + installNodeTopLevelExportPatches, + type NodeTopLevelExportPatchRuntime, +} from "./top-level-export-patches-node"; +import { createHook as createImportInTheMiddleHook } from "../import-in-the-middle/create-hook.mjs"; +import type { TopLevelImportHook } from "./top-level-export-patches"; + +describe("installNodeTopLevelExportPatches", () => { + it("uses in-process registerHooks when synchronous hooks are supported", () => { + const events: string[] = []; + const importHook = { + initialize(data: unknown) { + events.push(`initialize:${JSON.stringify(data)}`); + return Promise.resolve(); + }, + async load() {}, + loadSync() {}, + async resolve() {}, + resolveSync() {}, + }; + + installNodeTopLevelExportPatches({ + asyncImportHookUrl: + "file:///braintrust/hook.mjs?braintrust-iitm-loader=true", + hooks: [fakeHook()], + registryImportUrl: "file:///braintrust/hook.mjs", + runtime: fakeRuntime(events, { + createImportHook(meta, options) { + events.push(`create:${meta.url}:${options.registerUrl}`); + return importHook; + }, + register() { + events.push("register"); + }, + registerHooks(hooks) { + events.push( + `registerHooks:${typeof hooks.resolve}:${typeof hooks.load}`, + ); + }, + supportsSyncHooks: () => true, + }), + }); + + expect(events).toEqual([ + "create:file:///braintrust/hook.mjs:file:///braintrust/hook.mjs", + 'initialize:{"include":["pkg"]}', + "registerHooks:function:function", + "importHook:pkg", + "importHookPatched:patched", + "requireHook:pkg", + "requireHookPatched:patched", + ]); + }); + + it("falls back to self-registering hook.mjs as the async loader", () => { + const events: string[] = []; + const asyncImportHookUrl = + "file:///braintrust/hook.mjs?braintrust-iitm-loader=true"; + + installNodeTopLevelExportPatches({ + asyncImportHookUrl, + hooks: [fakeHook()], + registryImportUrl: "file:///braintrust/hook.mjs", + runtime: fakeRuntime(events, { + register(specifier, options) { + events.push(`register:${specifier}:${JSON.stringify(options)}`); + }, + registerHooks() { + events.push("registerHooks"); + }, + supportsSyncHooks: () => false, + }), + }); + + expect(events).toEqual([ + `register:${asyncImportHookUrl}:{"data":{"include":["pkg"]}}`, + "importHook:pkg", + "importHookPatched:patched", + "requireHook:pkg", + "requireHookPatched:patched", + ]); + }); +}); + +describe("merged import-in-the-middle wrapper generation", () => { + it("imports register from the canonical Braintrust hook URL", async () => { + const hook = createImportInTheMiddleHook( + { url: "file:///braintrust/hook.mjs?braintrust-iitm-loader=true" }, + { registerUrl: "file:///braintrust/hook.mjs" }, + ); + await hook.initialize({ include: ["target"] }); + + const resolved = (await hook.resolve( + "target", + { conditions: ["import"], parentURL: "file:///app.mjs" }, + async () => ({ format: "module", url: "file:///target.mjs" }), + )) as { url: string }; + const loaded = (await hook.load( + resolved.url, + { format: "module" }, + async () => ({ + format: "module", + source: 'export const value = "original";', + }), + )) as { source: string }; + + expect(loaded.source).toContain( + "import { register } from 'file:///braintrust/hook.mjs'", + ); + expect(loaded.source).not.toContain("lib/register.js"); + }); +}); + +function fakeHook(): TopLevelImportHook { + return { + hook(exportsValue) { + (exportsValue as { value: string }).value = "patched"; + }, + integrations: ["mastra"], + specifiers: ["pkg"], + targets: ["node"], + }; +} + +function fakeRuntime( + events: string[], + { + createImportHook, + register, + registerHooks, + supportsSyncHooks, + }: { + createImportHook?: NodeTopLevelExportPatchRuntime["createImportHook"]; + register?: (specifier: string, options?: unknown) => void; + registerHooks?: (hooks: { load: unknown; resolve: unknown }) => void; + supportsSyncHooks: () => boolean; + }, +): Partial { + class FakeImportHook { + constructor( + modules: string[], + hookFn: (exportsValue: unknown, name: string) => unknown, + ) { + events.push(`importHook:${modules.join(",")}`); + const namespace = { value: "original" }; + hookFn(namespace, "pkg"); + events.push(`importHookPatched:${namespace.value}`); + } + } + class FakeRequireHook { + constructor( + modules: string[], + hookFn: (exportsValue: unknown, name: string) => unknown, + ) { + events.push(`requireHook:${modules.join(",")}`); + const namespace = { value: "original" }; + hookFn(namespace, "pkg"); + events.push(`requireHookPatched:${namespace.value}`); + } + } + return { + createImportHook: + createImportHook ?? + (() => ({ + initialize: () => Promise.resolve(), + async load() {}, + loadSync() {}, + async resolve() {}, + resolveSync() {}, + })), + importHookConstructor: FakeImportHook, + moduleApi: { + register: + register ?? + (() => { + events.push("register"); + }), + registerHooks, + }, + requireHookConstructor: FakeRequireHook, + supportsSyncHooks, + }; +} diff --git a/js/src/auto-instrumentations/loader/top-level-export-patches-node.ts b/js/src/auto-instrumentations/loader/top-level-export-patches-node.ts new file mode 100644 index 000000000..d7d262fea --- /dev/null +++ b/js/src/auto-instrumentations/loader/top-level-export-patches-node.ts @@ -0,0 +1,108 @@ +import * as module from "node:module"; +import { debugLogger } from "../../debug-logger"; +import { + createHook as createImportInTheMiddleHook, + supportsSyncHooks, + type ImportInTheMiddleHook as ImportInTheMiddleHookApi, +} from "../import-in-the-middle/create-hook.mjs"; +import ImportInTheMiddleRuntimeHook from "./import-in-the-middle-runtime.js"; +import RequireInTheMiddleHook from "./require-in-the-middle-runtime.js"; +import { + getTopLevelImportHookSpecifiers, + runTopLevelImportHooks, + type TopLevelImportHook, +} from "./top-level-export-patches.js"; + +export interface InstallNodeTopLevelExportPatchesOptions { + hooks: readonly TopLevelImportHook[]; + asyncImportHookUrl: string; + registryImportUrl: string; + runtime?: Partial; +} + +type HookCallback = ( + exportsValue: unknown, + name: string, + baseDir?: string, +) => unknown; + +type HookConstructor = new (modules: string[], hookFn: HookCallback) => unknown; + +interface NodeModuleApi { + register: (specifier: string, options?: unknown) => void; + registerHooks?: (hooks: { load: unknown; resolve: unknown }) => void; +} + +export interface NodeTopLevelExportPatchRuntime { + createImportHook: ( + meta: { url: string }, + options: { registerUrl: string }, + ) => ImportInTheMiddleHookApi; + importHookConstructor: HookConstructor; + moduleApi: NodeModuleApi; + requireHookConstructor: HookConstructor; + supportsSyncHooks: () => boolean; +} + +export function installNodeTopLevelExportPatches({ + asyncImportHookUrl, + hooks, + registryImportUrl, + runtime, +}: InstallNodeTopLevelExportPatchesOptions): void { + const specifiers = getTopLevelImportHookSpecifiers(hooks); + if (specifiers.length === 0) return; + + const effectiveRuntime = getRuntime(runtime); + const hookCallback: HookCallback = (exportsValue, name, baseDir) => + runTopLevelImportHooks(hooks, exportsValue, { + baseDir, + moduleName: name, + }); + + try { + if ( + effectiveRuntime.moduleApi.registerHooks && + effectiveRuntime.supportsSyncHooks() + ) { + const importHook = effectiveRuntime.createImportHook( + { url: registryImportUrl }, + { registerUrl: registryImportUrl }, + ); + void importHook.initialize({ include: specifiers }); + effectiveRuntime.moduleApi.registerHooks({ + load: importHook.loadSync, + resolve: importHook.resolveSync, + }); + } else { + effectiveRuntime.moduleApi.register(asyncImportHookUrl, { + data: { include: specifiers }, + }); + } + + new effectiveRuntime.importHookConstructor(specifiers, hookCallback); + } catch (err) { + debugLogger.warn("Failed to install ESM top-level import hooks", err); + } + + try { + new effectiveRuntime.requireHookConstructor(specifiers, hookCallback); + } catch (err) { + debugLogger.warn("Failed to install CJS top-level import hooks", err); + } +} + +function getRuntime( + overrides: Partial | undefined, +): NodeTopLevelExportPatchRuntime { + return { + createImportHook: createImportInTheMiddleHook, + importHookConstructor: + ImportInTheMiddleRuntimeHook as unknown as HookConstructor, + moduleApi: module as unknown as NodeModuleApi, + requireHookConstructor: + RequireInTheMiddleHook as unknown as HookConstructor, + supportsSyncHooks, + ...overrides, + }; +} diff --git a/js/src/auto-instrumentations/loader/top-level-export-patches.test.ts b/js/src/auto-instrumentations/loader/top-level-export-patches.test.ts new file mode 100644 index 000000000..f9e8db3e2 --- /dev/null +++ b/js/src/auto-instrumentations/loader/top-level-export-patches.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, it } from "vitest"; +import { + buildTopLevelImportHookSourceWrapper, + getDefaultTopLevelImportHooks, + getTopLevelImportHookSpecifiers, + runTopLevelImportHooks, + type TopLevelImportHook, +} from "./top-level-export-patches"; + +describe("top-level import hook registry", () => { + it("returns the Mastra runtime specifiers for node targets", () => { + const hooks = getDefaultTopLevelImportHooks({ target: "node" }); + + expect(getTopLevelImportHookSpecifiers(hooks)).toEqual([ + "@mastra/core", + "@mastra/core/mastra", + "@mastra/observability", + ]); + }); + + it("filters hooks by disabled integration config", () => { + const hooks = getDefaultTopLevelImportHooks({ + disabledIntegrationConfig: { mastra: false }, + target: "node", + }); + + expect(hooks).toEqual([]); + }); + + it("does not include Mastra for browser targets", () => { + const hooks = getDefaultTopLevelImportHooks({ target: "browser" }); + + expect(hooks).toEqual([]); + }); + + it("runs hook callbacks against a mutable namespace", () => { + const hooks: TopLevelImportHook[] = [ + { + hook(exportsValue) { + exportsValue.value = "patched"; + }, + integrations: ["mastra"], + specifiers: ["pkg"], + targets: ["node"], + }, + ]; + const namespace = { value: "original" }; + + const result = runTopLevelImportHooks(hooks, namespace, { + moduleName: "pkg", + }); + + expect(result).toBe(namespace); + expect(namespace.value).toBe("patched"); + }); + + it("uses a returned replacement namespace when a hook provides one", () => { + const replacement = { value: "replacement" }; + const hooks: TopLevelImportHook[] = [ + { + hook() { + return replacement; + }, + integrations: ["mastra"], + specifiers: ["pkg"], + targets: ["node"], + }, + ]; + + const result = runTopLevelImportHooks( + hooks, + { value: "original" }, + { moduleName: "pkg" }, + ); + + expect(result).toBe(replacement); + }); + + it("passes runtime resolution context through hook callbacks", () => { + let seen: + | { baseDir: string | undefined; resolutionBase: string | undefined } + | undefined; + const hooks: TopLevelImportHook[] = [ + { + hook(_exportsValue, _name, baseDir, resolutionBase) { + seen = { baseDir, resolutionBase }; + }, + integrations: ["mastra"], + specifiers: ["pkg"], + targets: ["node"], + }, + ]; + + runTopLevelImportHooks( + hooks, + { value: "original" }, + { + baseDir: "/tmp/pkg", + moduleName: "pkg", + resolutionBase: "file:///tmp/app/bundle.mjs", + }, + ); + + expect(seen).toEqual({ + baseDir: "/tmp/pkg", + resolutionBase: "file:///tmp/app/bundle.mjs", + }); + }); + + it("generates an ESM wrapper that calls the shared hook runner", () => { + const wrapper = buildTopLevelImportHookSourceWrapper([fakeHook("node")], { + format: "esm", + modulePath: "index.js", + originalModuleSpecifier: "braintrust-top-level-original:0", + packageName: "pkg", + source: `export { value } from "./value.js";`, + target: "node", + }); + + expect(wrapper).toContain( + `import * as __braintrustOriginal from "braintrust-top-level-original:0"`, + ); + expect(wrapper).toContain("__braintrustTopLevelImportHookRunner"); + expect(wrapper).toContain( + `export * from "braintrust-top-level-original:0"`, + ); + expect(wrapper).toContain("import.meta.url"); + expect(wrapper).toContain(" as value"); + }); + + it("generates a CJS wrapper that calls the shared hook runner", () => { + const wrapper = buildTopLevelImportHookSourceWrapper([fakeHook("node")], { + format: "cjs", + modulePath: "index.cjs", + originalModuleSpecifier: "/tmp/pkg/index.cjs?braintrust-original", + packageName: "pkg", + source: `exports.value = "original";`, + target: "node", + }); + + expect(wrapper).toContain( + `const __braintrustOriginal = require("/tmp/pkg/index.cjs?braintrust-original")`, + ); + expect(wrapper).toContain("__braintrustTopLevelImportHookRunner"); + expect(wrapper).toContain("typeof __filename"); + expect(wrapper).toContain(`Object.defineProperty(exports, "value"`); + }); + + it("supports browser-safe hook descriptors while Mastra remains node-only", () => { + const wrapper = buildTopLevelImportHookSourceWrapper( + [fakeHook("browser")], + { + format: "esm", + modulePath: "index.js", + originalModuleSpecifier: "braintrust-top-level-original:1", + packageName: "pkg", + source: `export const value = "original";`, + target: "browser", + }, + ); + + expect(wrapper).toContain("__braintrustTopLevelImportHookRunner"); + expect(getDefaultTopLevelImportHooks({ target: "browser" })).toEqual([]); + }); +}); + +function fakeHook(target: "node" | "browser"): TopLevelImportHook { + return { + hook(exportsValue) { + exportsValue.value = "patched"; + }, + integrations: ["mastra"], + sourceTargets: [ + { + exportNames: ["value"], + modulePaths: + target === "node" ? ["index.js", "index.cjs"] : ["index.js"], + packageName: "pkg", + specifier: "pkg", + }, + ], + specifiers: ["pkg"], + targets: [target], + }; +} diff --git a/js/src/auto-instrumentations/loader/top-level-export-patches.ts b/js/src/auto-instrumentations/loader/top-level-export-patches.ts new file mode 100644 index 000000000..76a21be69 --- /dev/null +++ b/js/src/auto-instrumentations/loader/top-level-export-patches.ts @@ -0,0 +1,456 @@ +import type { InstrumentationIntegrationsConfig } from "../../instrumentation/config"; +import { isInstrumentationIntegrationDisabled } from "../../instrumentation/config"; +import { patchMastraExports } from "./mastra-observability-patch.js"; + +export type TopLevelImportHookTarget = "node" | "browser"; +export type TopLevelImportHookFormat = "esm" | "cjs"; +export type MutableExportNamespace = Record; + +export interface TopLevelImportHookSourceTarget { + packageName: string; + modulePaths: readonly string[]; + specifier: string; + exportNames: readonly string[]; +} + +export interface TopLevelImportHook { + integrations: readonly (keyof InstrumentationIntegrationsConfig)[]; + specifiers: readonly string[]; + targets: readonly TopLevelImportHookTarget[]; + sourceTargets?: readonly TopLevelImportHookSourceTarget[]; + hook( + exports: MutableExportNamespace, + name: string, + baseDir?: string, + resolutionBase?: string, + ): unknown | void; +} + +export interface TopLevelImportHookContext { + moduleName: string; + baseDir?: string; + resolutionBase?: string; +} + +export interface TopLevelImportHookSourceWrapperInput { + baseDir?: string; + format: TopLevelImportHookFormat; + modulePath: string; + originalModuleSpecifier: string; + packageName: string; + source: string; + target: TopLevelImportHookTarget; +} + +const TOP_LEVEL_IMPORT_HOOK_RUNNER_GLOBAL = + "__braintrustTopLevelImportHookRunner"; + +export function getDefaultTopLevelImportHooks({ + disabledIntegrationConfig, + target, +}: { + disabledIntegrationConfig?: InstrumentationIntegrationsConfig; + target: TopLevelImportHookTarget; +}): TopLevelImportHook[] { + return defaultTopLevelImportHooks.filter( + (hook) => + hook.targets.includes(target) && + !isInstrumentationIntegrationDisabled( + disabledIntegrationConfig, + ...hook.integrations, + ), + ); +} + +export function getTopLevelImportHookSpecifiers( + hooks: readonly TopLevelImportHook[], +): string[] { + return [...new Set(hooks.flatMap((hook) => hook.specifiers))]; +} + +export function runTopLevelImportHooks( + hooks: readonly TopLevelImportHook[], + exportsValue: unknown, + context: TopLevelImportHookContext, +): unknown { + let nextExports = exportsValue; + for (const hook of hooks) { + if (!matchesRuntimeHook(hook, context.moduleName)) { + continue; + } + + try { + const patched = hook.hook( + asMutableNamespace(nextExports), + context.moduleName, + context.baseDir, + context.resolutionBase, + ); + if (patched !== undefined) { + nextExports = patched; + } + } catch { + // Hook failures must never escape into user module evaluation. + } + } + return nextExports; +} + +export function installTopLevelImportHookRunner( + hooks: readonly TopLevelImportHook[], +): void { + Object.defineProperty(globalThis, TOP_LEVEL_IMPORT_HOOK_RUNNER_GLOBAL, { + configurable: true, + enumerable: false, + value( + exportsValue: unknown, + name: string, + baseDir?: string, + resolutionBase?: string, + ) { + return runTopLevelImportHooks(hooks, exportsValue, { + baseDir, + moduleName: name, + resolutionBase, + }); + }, + writable: true, + }); +} + +export function buildTopLevelImportHookSourceWrapper( + hooks: readonly TopLevelImportHook[], + input: TopLevelImportHookSourceWrapperInput, +): string | null { + const sourceTargets = hooks + .filter((hook) => hook.targets.includes(input.target)) + .flatMap((hook) => hook.sourceTargets ?? []) + .filter( + (sourceTarget) => + sourceTarget.packageName === input.packageName && + sourceTarget.modulePaths.includes(input.modulePath), + ); + + if (sourceTargets.length === 0) { + return null; + } + + const specifier = sourceTargets[0].specifier; + const exportNames = [ + ...new Set([ + ...sourceTargets.flatMap((target) => target.exportNames), + ...collectStaticExportNames(input.source, input.format), + ]), + ].sort((a, b) => Number(a === "default") - Number(b === "default")); + + if (exportNames.length === 0) { + return null; + } + + return input.format === "esm" + ? buildEsmSourceWrapper({ + baseDir: input.baseDir, + exportNames, + moduleName: specifier, + originalModuleSpecifier: input.originalModuleSpecifier, + }) + : buildCjsSourceWrapper({ + baseDir: input.baseDir, + exportNames, + moduleName: specifier, + originalModuleSpecifier: input.originalModuleSpecifier, + }); +} + +export { + getDefaultTopLevelImportHooks as getDefaultTopLevelExportPatches, + getTopLevelImportHookSpecifiers as getTopLevelExportPatchSpecifiers, + runTopLevelImportHooks as applyTopLevelExportRuntimePatches, +}; + +export type { + TopLevelImportHook as TopLevelExportPatch, + TopLevelImportHookContext as TopLevelExportPatchContext, + TopLevelImportHookFormat as TopLevelExportPatchFormat, + TopLevelImportHookTarget as TopLevelExportPatchTarget, +}; + +function matchesRuntimeHook( + hook: TopLevelImportHook, + moduleName: string, +): boolean { + return hook.specifiers.includes(moduleName); +} + +function asMutableNamespace(exportsValue: unknown): MutableExportNamespace { + if (exportsValue && typeof exportsValue === "object") { + return exportsValue as MutableExportNamespace; + } + + return { default: exportsValue }; +} + +function collectStaticExportNames( + source: string, + format: TopLevelImportHookFormat, +): string[] { + if (format === "cjs") { + const names = new Set(); + for (const match of source.matchAll( + /\bexports\.([A-Za-z_$][\w$]*)\s*=|\bmodule\.exports\.([A-Za-z_$][\w$]*)\s*=|Object\.defineProperty\s*\(\s*exports\s*,\s*["']([^"']+)["']/g, + )) { + names.add(match[1] ?? match[2] ?? match[3]); + } + return [...names]; + } + + const names = new Set(); + + for (const match of source.matchAll( + /\bexport\s+(?:async\s+)?(?:class|function|const|let|var)\s+([A-Za-z_$][\w$]*)/g, + )) { + names.add(match[1]); + } + + if (/\bexport\s+default\b/.test(source)) { + names.add("default"); + } + + for (const match of source.matchAll( + /\bexport\s*\{([^}]+)\}(?:\s*from\s*["'][^"']+["'])?/g, + )) { + for (const part of match[1].split(",")) { + const trimmed = part.trim(); + if (!trimmed || trimmed.startsWith("type ")) continue; + const aliasMatch = trimmed.match(/\bas\s+([A-Za-z_$][\w$]*|default)$/); + const directMatch = trimmed.match(/^([A-Za-z_$][\w$]*|default)$/); + const name = aliasMatch?.[1] ?? directMatch?.[1]; + if (name) names.add(name); + } + } + + return [...names]; +} + +function buildEsmSourceWrapper({ + baseDir, + exportNames, + moduleName, + originalModuleSpecifier, +}: { + baseDir?: string; + exportNames: readonly string[]; + moduleName: string; + originalModuleSpecifier: string; +}): string { + const locals = exportNames.map((name, index) => ({ + exportName: name, + localName: + name === "default" + ? "__braintrustDefaultExport" + : toSafeLocalName(name, index), + })); + + return `import * as __braintrustOriginal from ${JSON.stringify(originalModuleSpecifier)}; + +${locals + .map( + ({ exportName, localName }) => + `let ${localName} = __braintrustOriginal[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + +const __braintrustExports = Object.create(null); +${locals.map(({ exportName, localName }) => buildMutableExportDescriptor(exportName, localName)).join("\n")} + +try { + const __braintrustHookRunner = globalThis[${JSON.stringify(TOP_LEVEL_IMPORT_HOOK_RUNNER_GLOBAL)}]; + if (typeof __braintrustHookRunner === "function") { + const __braintrustPatched = __braintrustHookRunner( + __braintrustExports, + ${JSON.stringify(moduleName)}, + ${JSON.stringify(baseDir)}, + import.meta.url, + ); + if (__braintrustPatched && __braintrustPatched !== __braintrustExports) { +${locals + .map( + ({ exportName, localName }) => + ` if (Object.prototype.hasOwnProperty.call(__braintrustPatched, ${JSON.stringify(exportName)})) ${localName} = __braintrustPatched[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + } + } +} catch (e) { + // Top-level import hooks are best-effort in bundled output. +} + +export * from ${JSON.stringify(originalModuleSpecifier)}; +${buildEsmExports(locals)} +`; +} + +function buildCjsSourceWrapper({ + baseDir, + exportNames, + moduleName, + originalModuleSpecifier, +}: { + baseDir?: string; + exportNames: readonly string[]; + moduleName: string; + originalModuleSpecifier: string; +}): string { + const locals = exportNames.map((name, index) => ({ + exportName: name, + localName: + name === "default" + ? "__braintrustDefaultExport" + : toSafeLocalName(name, index), + })); + const cjsRequire = "require"; + + return `"use strict"; +const __braintrustOriginal = ${cjsRequire}(${JSON.stringify(originalModuleSpecifier)}); +const __braintrustPatchedExportNames = new Set(${JSON.stringify(exportNames)}); +try { + const __braintrustOriginalDescriptors = Object.getOwnPropertyDescriptors(__braintrustOriginal); +${locals.map(({ exportName }) => ` delete __braintrustOriginalDescriptors[${JSON.stringify(exportName)}];`).join("\n")} + Object.defineProperties(exports, __braintrustOriginalDescriptors); +} catch (e) { + for (const __braintrustKey in __braintrustOriginal) { + if (!__braintrustPatchedExportNames.has(__braintrustKey)) { + exports[__braintrustKey] = __braintrustOriginal[__braintrustKey]; + } + } +} + +${locals + .map( + ({ exportName, localName }) => + `let ${localName} = __braintrustOriginal[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + +const __braintrustExports = Object.create(null); +${locals.map(({ exportName, localName }) => buildMutableExportDescriptor(exportName, localName)).join("\n")} + +try { + const __braintrustHookRunner = globalThis[${JSON.stringify(TOP_LEVEL_IMPORT_HOOK_RUNNER_GLOBAL)}]; + if (typeof __braintrustHookRunner === "function") { + const __braintrustPatched = __braintrustHookRunner( + __braintrustExports, + ${JSON.stringify(moduleName)}, + ${JSON.stringify(baseDir)}, + typeof __filename === "string" ? __filename : undefined, + ); + if (__braintrustPatched && __braintrustPatched !== __braintrustExports) { +${locals + .map( + ({ exportName, localName }) => + ` if (Object.prototype.hasOwnProperty.call(__braintrustPatched, ${JSON.stringify(exportName)})) ${localName} = __braintrustPatched[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + } + } +} catch (e) { + // Top-level import hooks are best-effort in bundled output. +} + +${locals.map(({ exportName, localName }) => buildCjsExportDescriptor(exportName, localName)).join("\n")} +`; +} + +function buildMutableExportDescriptor( + exportName: string, + localName: string, +): string { + return `Object.defineProperty(__braintrustExports, ${JSON.stringify(exportName)}, { + configurable: true, + enumerable: true, + get() { return ${localName}; }, + set(value) { ${localName} = value; return true; }, +});`; +} + +function buildEsmExports( + locals: readonly { exportName: string; localName: string }[], +): string { + const named = locals.filter(({ exportName }) => exportName !== "default"); + const defaultExport = locals.find( + ({ exportName }) => exportName === "default", + ); + const lines = named.length + ? [ + `export { ${named.map(({ localName, exportName }) => (localName === exportName ? exportName : `${localName} as ${exportName}`)).join(", ")} };`, + ] + : []; + if (defaultExport) { + lines.push(`export { ${defaultExport.localName} as default };`); + } + return lines.join("\n"); +} + +function buildCjsExportDescriptor( + exportName: string, + localName: string, +): string { + return `Object.defineProperty(exports, ${JSON.stringify(exportName)}, { + configurable: true, + enumerable: true, + get() { return ${localName}; }, +});`; +} + +const MASTRA_CORE_ENTRY_PATHS = [ + "dist/index.js", + "dist/index.cjs", + "dist/mastra/index.js", + "dist/mastra/index.cjs", +]; +const MASTRA_OBSERVABILITY_ENTRY_PATHS = ["dist/index.js", "dist/index.cjs"]; + +const mastraTopLevelImportHook: TopLevelImportHook = { + hook(exportsValue, name, baseDir, resolutionBase) { + return patchMastraExports(exportsValue, { + baseDir, + moduleName: name, + resolutionBase, + }); + }, + integrations: ["mastra"], + sourceTargets: [ + { + exportNames: ["Mastra"], + modulePaths: MASTRA_CORE_ENTRY_PATHS.filter((path) => + path.startsWith("dist/index."), + ), + packageName: "@mastra/core", + specifier: "@mastra/core", + }, + { + exportNames: ["Mastra"], + modulePaths: MASTRA_CORE_ENTRY_PATHS.filter((path) => + path.startsWith("dist/mastra/"), + ), + packageName: "@mastra/core", + specifier: "@mastra/core/mastra", + }, + { + exportNames: ["Observability"], + modulePaths: MASTRA_OBSERVABILITY_ENTRY_PATHS, + packageName: "@mastra/observability", + specifier: "@mastra/observability", + }, + ], + specifiers: ["@mastra/core", "@mastra/core/mastra", "@mastra/observability"], + targets: ["node"], +}; + +const defaultTopLevelImportHooks: readonly TopLevelImportHook[] = [ + mastraTopLevelImportHook, +]; + +function toSafeLocalName(exportName: string, index: number): string { + return `__braintrust_export_${index}_${exportName.replace(/\W/g, "_")}`; +} diff --git a/js/src/auto-instrumentations/require-in-the-middle/index.d.ts b/js/src/auto-instrumentations/require-in-the-middle/index.d.ts new file mode 100644 index 000000000..71421bb11 --- /dev/null +++ b/js/src/auto-instrumentations/require-in-the-middle/index.d.ts @@ -0,0 +1,8 @@ +export type OnRequireFn = (exports: T, name: string, basedir?: string) => T; + +export declare class Hook { + constructor(modules: string[], onrequire: OnRequireFn); + unhook(): void; +} + +export default Hook; diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index a1fe0e352..dc6c51e29 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -202,9 +202,8 @@ export class BraintrustPlugin extends BasePlugin { // Mastra is intentionally not wired here: `@mastra/core` ships its own // ObservabilityExporter contract, and `BraintrustObservabilityExporter` - // (wrappers/mastra.ts) is auto-installed by the loader patch in - // `auto-instrumentations/loader/mastra-observability-patch.ts` rather than - // by a BasePlugin / tracingChannel subscription. + // (wrappers/mastra.ts) is auto-installed by the top-level import hook + // registry rather than by a BasePlugin / tracingChannel subscription. } protected onDisable(): void { diff --git a/js/src/node/apply-auto-instrumentation-entry.ts b/js/src/node/apply-auto-instrumentation-entry.ts index 0328688dd..d30bb43d3 100644 --- a/js/src/node/apply-auto-instrumentation-entry.ts +++ b/js/src/node/apply-auto-instrumentation-entry.ts @@ -3,7 +3,18 @@ import { register } from "node:module"; import { pathToFileURL } from "node:url"; import { getDefaultAutoInstrumentationConfigs } from "../auto-instrumentations/configs/all"; import { ModulePatch } from "../auto-instrumentations/loader/cjs-patch"; +import { installMastraExporterFactory } from "../auto-instrumentations/loader/mastra-observability-patch"; +import { + getDefaultTopLevelImportHooks, + installTopLevelImportHookRunner, +} from "../auto-instrumentations/loader/top-level-export-patches"; +import { installNodeTopLevelExportPatches } from "../auto-instrumentations/loader/top-level-export-patches-node"; import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel"; +import { + isInstrumentationIntegrationDisabled, + readDisabledInstrumentationEnvConfig, +} from "../instrumentation/config"; +import { BraintrustObservabilityExporter } from "../wrappers/mastra"; interface ApplyAutoInstrumentationState { applied?: boolean; @@ -33,8 +44,29 @@ if (!state.applied) { patchTracingChannel(diagnostics_channel.tracingChannel); const allConfigs = getDefaultAutoInstrumentationConfigs(); + const disabled = readDisabledInstrumentationEnvConfig( + process.env.BRAINTRUST_DISABLE_INSTRUMENTATION, + ).integrations; + if (!isInstrumentationIntegrationDisabled(disabled, "mastra")) { + installMastraExporterFactory(() => new BraintrustObservabilityExporter()); + } const currentModuleUrl = getCurrentModuleUrl(); + const topLevelImportHooks = getDefaultTopLevelImportHooks({ + disabledIntegrationConfig: disabled, + target: "node", + }); + const autoInstrumentationHookUrl = + getAutoInstrumentationHookUrl(currentModuleUrl); + installTopLevelImportHookRunner(topLevelImportHooks); + installNodeTopLevelExportPatches({ + asyncImportHookUrl: getImportInTheMiddleLoaderUrl( + autoInstrumentationHookUrl, + ), + hooks: topLevelImportHooks, + registryImportUrl: autoInstrumentationHookUrl, + }); + register("./auto-instrumentations/loader/esm-hook.mjs", { parentURL: currentModuleUrl, data: { instrumentations: allConfigs }, @@ -71,4 +103,14 @@ function getCurrentModuleUrl(): string { return pathToFileURL(process.argv[1] ?? process.cwd()).href; } +function getAutoInstrumentationHookUrl(currentModuleUrl: string): string { + return new URL("./auto-instrumentations/hook.mjs", currentModuleUrl).href; +} + +function getImportInTheMiddleLoaderUrl(canonicalUrl: string): string { + const parsed = new URL(canonicalUrl); + parsed.searchParams.set("braintrust-iitm-loader", "true"); + return parsed.href; +} + export {}; diff --git a/js/src/node/config.ts b/js/src/node/config.ts index 09c27b807..cb5e7a034 100644 --- a/js/src/node/config.ts +++ b/js/src/node/config.ts @@ -20,6 +20,10 @@ import { readDisabledInstrumentationEnvConfig, } from "../instrumentation/config"; import { installMastraExporterFactory } from "../auto-instrumentations/loader/mastra-observability-patch"; +import { + getDefaultTopLevelImportHooks, + installTopLevelImportHookRunner, +} from "../auto-instrumentations/loader/top-level-export-patches"; import { BraintrustObservabilityExporter } from "../wrappers/mastra"; const BRAINTRUST_ENV_SEARCH_PARENT_LIMIT = 64; @@ -136,15 +140,19 @@ export function configureNode() { _internalSetInitialState(); - // The Mastra patches rewritten by the bundler plugin and the loader hook - // both look up an exporter factory on `globalThis` at runtime. Register it - // here too so a bundled app (where neither hook.mjs nor the loader runs - // against `@mastra/core`) still gets the exporter installed when its code - // imports `braintrust`. The loader path (hook.mjs) also calls this; the - // `??=` makes double-registration a no-op. + // Bundled wrapper modules call the top-level import hook runner through + // `globalThis`, and Mastra's hook uses the exporter factory below. Register + // both when the app imports `braintrust`; the loader path (hook.mjs) does + // the same setup for zero-line instrumentation. const disabled = readDisabledInstrumentationEnvConfig( iso.getEnv("BRAINTRUST_DISABLE_INSTRUMENTATION"), ).integrations; + installTopLevelImportHookRunner( + getDefaultTopLevelImportHooks({ + disabledIntegrationConfig: disabled, + target: "node", + }), + ); if (!isInstrumentationIntegrationDisabled(disabled, "mastra")) { installMastraExporterFactory(() => new BraintrustObservabilityExporter()); } diff --git a/js/src/wrappers/mastra.ts b/js/src/wrappers/mastra.ts index 9946a9643..6542c1cc3 100644 --- a/js/src/wrappers/mastra.ts +++ b/js/src/wrappers/mastra.ts @@ -10,9 +10,9 @@ * Two integration paths: * - **Manual**: `new Mastra({ observability: new Observability({ configs: { * default: { exporters: [new BraintrustObservabilityExporter()] } } }) })` - * - **Auto** (under `node --import braintrust/hook.mjs`): the loader patches - * `@mastra/core`'s `dist/mastra/index.{js,cjs}` to wrap `Mastra` so it - * calls `defaultInstance.registerExporter(exporter)` after construction. + * - **Auto** (under `node --import braintrust/hook.mjs`): the top-level + * import hook registry wraps Mastra's exported constructors so a default + * Observability config receives this exporter automatically. * * Minimum supported Mastra version: 1.20.0 (when `Mastra.prototype.register` * `Exporter` and `ObservabilityInstance.registerExporter` were added). The @@ -370,7 +370,7 @@ function logExporterError(err: unknown): void { * * - **Auto-instrumentation**: run your app with * `node --import braintrust/hook.mjs`. The loader installs - * `BraintrustObservabilityExporter` into every `new Mastra(...)` + * `BraintrustObservabilityExporter` into Mastra Observability configs * automatically. * - **Manual wiring**: pass the exporter yourself: * diff --git a/js/tests/auto-instrumentations/fixtures/import-hook-query-mode.mjs b/js/tests/auto-instrumentations/fixtures/import-hook-query-mode.mjs new file mode 100644 index 000000000..3ad020cf0 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/import-hook-query-mode.mjs @@ -0,0 +1,15 @@ +import { parentPort } from "node:worker_threads"; + +const imported = await import(process.env.BRAINTRUST_QUERY_HOOK_URL); +const state = globalThis[Symbol.for("braintrust.applyAutoInstrumentation")]; + +parentPort?.postMessage({ + result: { + applied: state?.applied === true, + hasInitialize: typeof imported.initialize === "function", + hasLoad: typeof imported.load === "function", + hasRegister: typeof imported.register === "function", + hasResolve: typeof imported.resolve === "function", + }, + type: "hook-query-result", +}); diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/package.json b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/package.json new file mode 100644 index 000000000..84b0d6964 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/package.json @@ -0,0 +1,15 @@ +{ + "name": "@mastra/core", + "version": "1.26.0", + "type": "module", + "exports": { + ".": { + "import": "./dist/index.js", + "require": "./dist/index.cjs" + }, + "./mastra": { + "import": "./dist/mastra/index.js", + "require": "./dist/mastra/index.cjs" + } + } +} diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json new file mode 100644 index 000000000..ed647752e --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json @@ -0,0 +1,11 @@ +{ + "name": "@mastra/observability", + "version": "1.26.0", + "type": "module", + "exports": { + ".": { + "import": "./dist/index.js", + "require": "./dist/index.cjs" + } + } +} diff --git a/js/tests/auto-instrumentations/fixtures/runtime-apply-auto-mastra-top-level-esm.mjs b/js/tests/auto-instrumentations/fixtures/runtime-apply-auto-mastra-top-level-esm.mjs new file mode 100644 index 000000000..94c9ab6c0 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/runtime-apply-auto-mastra-top-level-esm.mjs @@ -0,0 +1,3 @@ +import "braintrust/apply-auto-instrumentation"; + +await import("./test-mastra-top-level-esm.mjs"); diff --git a/js/tests/auto-instrumentations/fixtures/test-mastra-bundler.js b/js/tests/auto-instrumentations/fixtures/test-mastra-bundler.js new file mode 100644 index 000000000..26a50077c --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/test-mastra-bundler.js @@ -0,0 +1,3 @@ +import { Mastra } from "@mastra/core"; + +new Mastra({}); diff --git a/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-cjs.cjs b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-cjs.cjs new file mode 100644 index 000000000..0e761531f --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-cjs.cjs @@ -0,0 +1,14 @@ +const { parentPort } = require("node:worker_threads"); +const { Mastra } = require("@mastra/core"); + +const mastra = new Mastra({}); +parentPort?.postMessage({ + result: { + exporters: + mastra.observability?.config?.configs?.default?.exporters?.map( + (exporter) => exporter.name, + ) ?? [], + hasObservability: Boolean(mastra.observability), + }, + type: "mastra-result", +}); diff --git a/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-esm.mjs b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-esm.mjs new file mode 100644 index 000000000..e60a93761 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-esm.mjs @@ -0,0 +1,14 @@ +import { parentPort } from "node:worker_threads"; +import { Mastra } from "@mastra/core"; + +const mastra = new Mastra({}); +parentPort?.postMessage({ + result: { + exporters: + mastra.observability?.config?.configs?.default?.exporters?.map( + (exporter) => exporter.name, + ) ?? [], + hasObservability: Boolean(mastra.observability), + }, + type: "mastra-result", +}); diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-app.mjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-app.mjs index e6d69f304..b56411050 100644 --- a/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-app.mjs +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-app.mjs @@ -2,9 +2,11 @@ import assert from "node:assert"; import getLabel, { foo } from "hook-target"; import getUnhookedLabel, { foo as unhookedFoo } from "unhooked-target"; import cjsTarget from "cjs-hook-target"; +import { nestedValue } from "cjs-reexport-target"; assert.equal(foo, 57); assert.equal(getLabel(), "patched"); assert.equal(unhookedFoo, 10); assert.equal(getUnhookedLabel(), "untouched"); assert.equal(cjsTarget.value, 8); +assert.equal(nestedValue, "nested"); diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-setup.mjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-setup.mjs index 9e7377fed..abbd89c92 100644 --- a/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-setup.mjs +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-async-setup.mjs @@ -1,5 +1,6 @@ import assert from "node:assert"; import { register } from "node:module"; +import { fileURLToPath } from "node:url"; import { Hook, createAddHookMessageChannel, @@ -12,24 +13,27 @@ const hookUrl = new URL( const { addHookMessagePort, registerOptions, waitForAllMessagesAcknowledged } = createAddHookMessageChannel(); -register(hookUrl.href, import.meta.url, registerOptions); +register(fileURLToPath(hookUrl), import.meta.url, registerOptions); globalThis.__braintrustIitmAsyncHookCalls = 0; -const hook = new Hook(["hook-target", "cjs-hook-target"], (exports, name) => { - globalThis.__braintrustIitmAsyncHookCalls++; - if (name === "hook-target") { - exports.foo += 15; - exports.default = () => "patched"; - } - if (name === "cjs-hook-target") { - exports.default.value = 8; - } -}); +const hook = new Hook( + ["hook-target", "cjs-hook-target", "cjs-reexport-target"], + (exports, name) => { + globalThis.__braintrustIitmAsyncHookCalls++; + if (name === "hook-target") { + exports.foo += 15; + exports.default = () => "patched"; + } + if (name === "cjs-hook-target") { + exports.default.value = 8; + } + }, +); await waitForAllMessagesAcknowledged(); process.on("exit", () => { - assert.equal(globalThis.__braintrustIitmAsyncHookCalls, 2); + assert.equal(globalThis.__braintrustIitmAsyncHookCalls, 3); hook.unhook(); addHookMessagePort.close(); }); diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-sync-app.mjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-sync-app.mjs index 9f4bb0012..9ee4c7878 100644 --- a/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-sync-app.mjs +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/iitm-sync-app.mjs @@ -14,7 +14,7 @@ register(); let calls = 0; const hook = new Hook( - ["hook-target", "cjs-hook-target", "fs"], + ["hook-target", "cjs-hook-target", "cjs-reexport-target", "fs"], (exports, name) => { calls++; if (name === "hook-target") { @@ -33,6 +33,7 @@ const hook = new Hook( const target = await import("hook-target"); const other = await import("unhooked-target"); const cjsTarget = await import("cjs-hook-target"); +const cjsReexportTarget = await import("cjs-reexport-target"); const fs = await import("node:fs"); assert.equal(target.foo, 57); @@ -40,11 +41,13 @@ assert.equal(target.default(), "patched"); assert.equal(other.foo, 10); assert.equal(other.default(), "untouched"); assert.equal(cjsTarget.default.value, 8); +assert.equal(cjsReexportTarget.nestedValue, "nested"); +assert.equal(cjsReexportTarget.rootValue, undefined); assert.equal(fs.existsSync("/definitely/not/a/real/path"), true); const require = createRequire(import.meta.url); const requiredFs = require("fs"); assert.equal(Object.isExtensible(requiredFs), true); -assert.equal(calls, 3); +assert.equal(calls, 4); hook.unhook(); diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-leaf/index.cjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-leaf/index.cjs new file mode 100644 index 000000000..a8b25ed8f --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-leaf/index.cjs @@ -0,0 +1 @@ +exports.rootValue = "root"; diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-leaf/package.json b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-leaf/package.json new file mode 100644 index 000000000..dbeb183ab --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-leaf/package.json @@ -0,0 +1,5 @@ +{ + "name": "cjs-reexport-leaf", + "version": "1.0.0", + "main": "./index.cjs" +} diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/index.cjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/index.cjs new file mode 100644 index 000000000..22730cdc1 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/index.cjs @@ -0,0 +1 @@ +module.exports = require("cjs-reexport-leaf"); diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/node_modules/cjs-reexport-leaf/index.cjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/node_modules/cjs-reexport-leaf/index.cjs new file mode 100644 index 000000000..673c5564d --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/node_modules/cjs-reexport-leaf/index.cjs @@ -0,0 +1 @@ +exports.nestedValue = "nested"; diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/node_modules/cjs-reexport-leaf/package.json b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/node_modules/cjs-reexport-leaf/package.json new file mode 100644 index 000000000..dbeb183ab --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/node_modules/cjs-reexport-leaf/package.json @@ -0,0 +1,5 @@ +{ + "name": "cjs-reexport-leaf", + "version": "1.0.0", + "main": "./index.cjs" +} diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/package.json b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/package.json new file mode 100644 index 000000000..3e6393aca --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/node_modules/cjs-reexport-target/package.json @@ -0,0 +1,8 @@ +{ + "name": "cjs-reexport-target", + "version": "1.0.0", + "main": "./index.cjs", + "exports": { + ".": "./index.cjs" + } +} diff --git a/js/tests/auto-instrumentations/loader-hook.test.ts b/js/tests/auto-instrumentations/loader-hook.test.ts index 51a20be41..bccae5dcd 100644 --- a/js/tests/auto-instrumentations/loader-hook.test.ts +++ b/js/tests/auto-instrumentations/loader-hook.test.ts @@ -20,6 +20,10 @@ const helperPromisePath = path.join( fixturesDir, "test-api-promise-preservation.mjs", ); +const importHookQueryModePath = path.join( + fixturesDir, + "import-hook-query-mode.mjs", +); const runtimeApplyAutoSideEffectEsmPath = path.join( fixturesDir, "runtime-apply-auto-side-effect-esm.mjs", @@ -28,6 +32,18 @@ const runtimeApplyAutoSideEffectCjsPath = path.join( fixturesDir, "runtime-apply-auto-side-effect-cjs.cjs", ); +const mastraTopLevelEsmPath = path.join( + fixturesDir, + "test-mastra-top-level-esm.mjs", +); +const mastraTopLevelCjsPath = path.join( + fixturesDir, + "test-mastra-top-level-cjs.cjs", +); +const runtimeApplyAutoMastraTopLevelEsmPath = path.join( + fixturesDir, + "runtime-apply-auto-mastra-top-level-esm.mjs", +); interface TestResult { events: { start: any[]; end: any[]; error: any[] }; @@ -83,6 +99,74 @@ describe("Unified Loader Hook Integration Tests", () => { expect(result.withResponseOk).toBe(true); expect(result.constructorName).toBe("HelperPromise"); }); + + it("should expose import hook exports in query mode without bootstrapping", async () => { + const result = await runWithWorkerMessage<{ + applied: boolean; + hasInitialize: boolean; + hasLoad: boolean; + hasRegister: boolean; + hasResolve: boolean; + }>({ + env: { + BRAINTRUST_QUERY_HOOK_URL: `${pathToFileURL(hookPath).href}?braintrust-iitm-loader=true`, + }, + execArgv: [], + messageType: "hook-query-result", + script: importHookQueryModePath, + }); + + expect(result).toEqual({ + applied: false, + hasInitialize: true, + hasLoad: true, + hasRegister: true, + hasResolve: true, + }); + }); + + it("should patch Mastra top-level ESM exports", async () => { + const result = await runWithWorkerMessage<{ + exporters: string[]; + hasObservability: boolean; + }>({ + execArgv: ["--import", hookPath], + messageType: "mastra-result", + script: mastraTopLevelEsmPath, + }); + + expect(result.hasObservability).toBe(true); + expect(result.exporters).toEqual(["braintrust"]); + }); + + it("should patch Mastra top-level CJS exports", async () => { + const result = await runWithWorkerMessage<{ + exporters: string[]; + hasObservability: boolean; + }>({ + execArgv: ["--import", hookPath], + messageType: "mastra-result", + script: mastraTopLevelCjsPath, + }); + + expect(result.hasObservability).toBe(true); + expect(result.exporters).toEqual(["braintrust"]); + }); + + it("should respect Mastra disable config for top-level export patches", async () => { + const result = await runWithWorkerMessage<{ + exporters: string[]; + hasObservability: boolean; + }>({ + env: { BRAINTRUST_DISABLE_INSTRUMENTATION: "mastra" }, + execArgv: ["--import", hookPath], + messageType: "mastra-result", + script: mastraTopLevelEsmPath, + }); + + expect(result.hasObservability).toBe(false); + expect(result.exporters).toEqual([]); + }); }); describe("apply-auto-instrumentation side-effect runtime setup", () => { @@ -126,6 +210,20 @@ describe("Unified Loader Hook Integration Tests", () => { expect(result.events.start.length).toBe(1); expect(result.events.end.length).toBe(1); }); + + it("should apply Mastra top-level export patches through the side-effect export", async () => { + const result = await runWithWorkerMessage<{ + exporters: string[]; + hasObservability: boolean; + }>({ + execArgv: [], + messageType: "mastra-result", + script: runtimeApplyAutoMastraTopLevelEsmPath, + }); + + expect(result.hasObservability).toBe(true); + expect(result.exporters).toEqual(["braintrust"]); + }); }); }); diff --git a/js/tests/auto-instrumentations/transformation.test.ts b/js/tests/auto-instrumentations/transformation.test.ts index fc402096b..31d82fec3 100644 --- a/js/tests/auto-instrumentations/transformation.test.ts +++ b/js/tests/auto-instrumentations/transformation.test.ts @@ -306,6 +306,42 @@ describe("Orchestrion Transformation Tests", () => { expect(output).not.toContain("TracingChannel"); }); + it("should apply node top-level export patches", async () => { + const { braintrustEsbuildPlugin } = + await import("../../src/auto-instrumentations/bundler/esbuild.js"); + + const entryPoint = path.join(fixturesDir, "test-mastra-bundler.js"); + const outfile = path.join(outputDir, "esbuild-mastra-bundle.js"); + + const result = await esbuild.build({ + absWorkingDir: fixturesDir, + bundle: true, + entryPoints: [entryPoint], + format: "esm", + logLevel: "error", + outfile, + platform: "node", + plugins: [braintrustEsbuildPlugin()], + preserveSymlinks: true, + write: true, + }); + + expect(result.errors).toHaveLength(0); + expect(fs.existsSync(outfile)).toBe(true); + + const output = fs.readFileSync(outfile, "utf-8"); + + expect(output).toContain("__braintrustTopLevelImportHookRunner"); + expect(output).toContain("__braintrustExports"); + expect(output.replace(/\\/g, "/")).not.toContain( + JSON.stringify( + path + .join(fixturesDir, "node_modules", "@mastra", "core") + .replace(/\\/g, "/"), + ), + ); + }); + it("should bundle dc-browser module when useDiagnosticChannelCompatShim is true", async () => { const { braintrustEsbuildPlugin } = await import("../../src/auto-instrumentations/bundler/esbuild.js"); @@ -342,6 +378,37 @@ describe("Orchestrion Transformation Tests", () => { // Should NOT import from external diagnostics_channel expect(output).not.toMatch(/from\s+["']diagnostics_channel["']/); }); + + it("should skip node-only top-level export patches for browser bundles", async () => { + const { braintrustEsbuildPlugin } = + await import("../../src/auto-instrumentations/bundler/esbuild.js"); + + const entryPoint = path.join(fixturesDir, "test-mastra-bundler.js"); + const outfile = path.join(outputDir, "esbuild-mastra-browser-bundle.js"); + + const result = await esbuild.build({ + absWorkingDir: fixturesDir, + bundle: true, + entryPoints: [entryPoint], + format: "esm", + logLevel: "error", + outfile, + platform: "browser", + plugins: [ + braintrustEsbuildPlugin({ useDiagnosticChannelCompatShim: true }), + ], + preserveSymlinks: true, + write: true, + }); + + expect(result.errors).toHaveLength(0); + expect(fs.existsSync(outfile)).toBe(true); + + const output = fs.readFileSync(outfile, "utf-8"); + + expect(output).not.toContain("__braintrustTopLevelImportHookRunner"); + expect(output).not.toContain("__braintrustOriginal"); + }); }); describe("vite", () => { @@ -570,6 +637,37 @@ describe("Orchestrion Transformation Tests", () => { expect(output).toContain("orchestrion:openai:chat.completions.create"); }); + it("should apply top-level export patches (turbopack loader-only mode)", async () => { + const { errors, output } = await runWebpackWithLoader({ + entry: path.join(fixturesDir, "test-mastra-bundler.js"), + output: { + path: outputDir, + filename: "turbopack-mastra-bundle.js", + library: { type: "module" }, + }, + experiments: { outputModule: true }, + mode: "development", + resolve: { modules: [nodeModulesDir, "node_modules"] }, + externals: { "node:module": "module node:module" }, + module: { + rules: [ + { + use: [ + { + loader: webpackLoaderPath, + options: { browser: false }, + }, + ], + }, + ], + }, + }); + + expect(errors).toHaveLength(0); + expect(output).toContain("__braintrustTopLevelImportHookRunner"); + expect(output).toContain("__braintrustExports"); + }); + it("should bundle dc-browser polyfill when browser: true (turbopack loader-only mode)", async () => { const { errors, output } = await runWebpackWithLoader({ entry: path.join(fixturesDir, "test-app.js"),